The Architecture Question: Can We Split Inference and Training Across GPUs?

In the middle of a complex machine learning pipeline design, a single question from the user cuts to the heart of the architectural challenge: "pro6000 - can we inference on first gpu (or two) and train on another one or two?"

This message, brief as it is, represents a pivotal moment in the DFlash training pipeline design. It is a response to the assistant's detailed breakdown of next steps and compute needs ([msg 7689]), which had laid out a sobering reality: the offline hidden state extraction approach was dead in the water due to a ~90 TB storage requirement, and the recommended alternative—online training where the target model and drafter run simultaneously on the same GPUs—faced severe memory constraints on the available hardware.

The Context That Produced This Question

To understand why this question matters, we must first understand the problem it aims to solve. The assistant's preceding message ([msg 7689]) had walked through the full pipeline for training a DFlash speculative decoding drafter for Qwen3.6-27B. The pipeline had three core phases: re-tokenizing 902,087 generated completions, extracting hidden states from the target model at five specific layers, and training the drafter using those hidden states.

The critical bottleneck was Phase 2. Each token in the dataset required storing hidden states from 5 layers × 5,120 hidden dimensions in BF16 format—51.2 KB per token. With completions averaging ~2,000 tokens each across 902K samples, the total storage requirement came to approximately 90 TB. This was not merely impractical; it was disqualifying. The team had already experienced a similar issue earlier when the old prompt-only dataset produced 645 GB of hidden states for only 143K samples before being abandoned.

The assistant's recommended solution was "Option A: Online Training"—load both the frozen target model (Qwen3.6-27B, ~54 GB in BF16) and the trainable drafter (~3.3 GB plus ~24 GB for AdamW optimizer states) on the same GPUs, extract hidden states during the forward pass, and feed them directly to the drafter without ever writing them to disk. This eliminated the storage problem entirely but introduced a memory problem: the combined model footprint of ~100-110 GB exceeded the 96 GB available on each PRO 6000 GPU.

The assistant's proposed workaround was tensor parallelism (TP=2), splitting the target model across two GPUs so each held ~27 GB of the model, leaving room for the drafter and optimizer on the same two GPUs. This would work, but it meant that all four GPUs in the PRO 6000 node would be tied up running the combined inference-plus-training workload, with no clear path to scaling throughput.

The User's Architectural Insight

The user's question proposes a fundamentally different architecture. Instead of running both inference and training on the same set of GPUs (whether TP=1 or TP=2), why not split the workload across GPUs? Use GPU 0 (or GPUs 0 and 1) exclusively for running the target model forward pass to extract hidden states, and use GPU 2 (or GPUs 2 and 3) exclusively for training the drafter on those extracted hidden states.

This is a data-parallel pipeline split rather than a model-parallel split. It is a heterogeneous GPU assignment where different GPUs in the same node serve different roles. The inference GPUs hold only the frozen target model and run forward passes, producing hidden states. The training GPUs hold only the drafter and optimizer, consuming those hidden states to compute gradients and update weights.

The question reveals several important assumptions and insights on the user's part. First, the user understands that the PRO 6000 node has four GPUs and is thinking about how to use them optimally rather than uniformly. Second, the user recognizes that the inference workload (forward pass through a 27B model) and the training workload (drafter forward + backward pass) have different memory profiles and might benefit from dedicated hardware. Third, the user is implicitly asking about the feasibility of GPU-to-GPU data transfer for the hidden states—can the output of one GPU's computation be efficiently consumed by another GPU in the same node?

The Technical Tradeoffs Embedded in the Question

The user's proposal introduces a critical variable that the assistant's TP=2 approach avoided: data movement. In the TP=2 approach, hidden states are produced and consumed within the same GPU (or across the TP group), requiring no inter-GPU transfer of activations beyond what the tensor-parallel communication already handles. In the user's proposed split, hidden states must be transferred from the inference GPUs to the training GPUs, likely over PCIe Gen5 or NVLink.

The PRO 6000 (RTX PRO 6000 Blackwell) GPUs are connected via NVLink on the same board, offering significantly higher bandwidth than PCIe. However, the hidden state tensors are substantial: for a batch of sequences with 512 anchor positions (as used in DFlash), each with 5 layers × 5,120 hidden dimensions in BF16, a single sample's hidden states would be approximately 26 MB. At a throughput of 50-100 samples per second, the inter-GPU bandwidth requirement would be in the range of 1.3-2.6 GB/s—well within NVLink capabilities but potentially a bottleneck if not managed carefully.

There is also the question of synchronization. In the TP=2 approach, inference and training happen in lockstep: the forward pass produces hidden states, and the drafter consumes them immediately within the same GPU operation. In the split approach, the inference GPUs produce hidden states that must be consumed by the training GPUs, introducing a producer-consumer synchronization pattern. This could be implemented as a pipelined data queue where inference GPUs push hidden state tensors to a shared buffer (e.g., in CPU memory or via GPU direct access) and training GPUs pull from it, but the latency and throughput characteristics would need careful tuning.

The Deeper Significance

What makes this question significant is not just its technical content but its timing and framing. It comes immediately after the assistant presented a detailed analysis that seemed to converge on a specific recommendation (TP=2 on PRO 6000 or renting B200s). The user's question opens up an alternative that the assistant had not considered: heterogeneous GPU roles within a single node.

This is a pattern that recurs throughout the conversation: the user frequently proposes architectural alternatives that the assistant has not enumerated, often based on a more intuitive understanding of the hardware topology. The assistant's analysis had framed the problem as a binary choice between "fit both models on one GPU" (impossible on PRO 6000) and "split the model across GPUs with TP" (possible but constraining). The user's question introduces a third option: split the workload, not the model.

The question also implicitly challenges the assistant's assumption that online training must mean "both models on the same GPU." The user recognizes that "online" can mean something broader: a pipeline where inference and training happen concurrently on different hardware, connected by a data transfer channel, rather than sequentially on the same hardware. This is a subtle but important reframing of the architecture.

The Knowledge Required to Understand This Question

To fully grasp the user's question, one needs to understand several pieces of context. The PRO 6000 node has four NVIDIA RTX PRO 6000 Blackwell GPUs, each with 96 GB of VRAM, connected by NVLink. The target model (Qwen3.6-27B) requires approximately 54 GB in BF16 precision. The DFlash drafter has approximately 2 billion parameters, requiring ~3.3 GB for the model weights and ~24 GB for AdamW optimizer states (parameters, gradients, and two momentum buffers). The hidden states that need to be transferred are from five specific layers (indices 1, 16, 31, 46, and 61) with a hidden dimension of 5,120.

One must also understand the distinction between tensor parallelism (splitting a single model's layers across multiple GPUs for a single forward pass) and data parallelism (running the same model on multiple GPUs with different data). The user's proposal is closer to a pipeline parallelism approach, where different stages of the computation run on different devices.

What This Question Produces

This question sets the stage for the next phase of the conversation. It forces a re-examination of the GPU topology and the data flow between inference and training. The assistant will need to analyze whether the NVLink bandwidth between GPUs on the PRO 6000 board is sufficient for the hidden state transfer, whether the producer-consumer synchronization can be implemented efficiently in PyTorch, and whether the overall throughput of the split approach exceeds the TP=2 approach.

More broadly, this question crystallizes the architectural decision that will determine the entire training strategy. The answer will determine whether the team can train on their existing PRO 6000 hardware (avoiding the ~$500-1,000 cost of renting B200 instances) and whether the training pipeline can achieve acceptable throughput. It is, in many ways, the central engineering question of this segment of the conversation.