The Feasibility Calculus: How a Single Python Script Decided the Fate of a 5-Day ML Training Run

Introduction

In the high-stakes world of large-scale machine learning, the difference between a successful training run and a catastrophic waste of compute resources often comes down to a single question: will this fit, and how long will it take? Experienced ML engineers know that launching a multi-day training job without a rigorous feasibility analysis is like setting sail without checking the weather. Message [msg 7754] of this opencode session captures that exact moment of reckoning — a detailed, back-of-the-envelope time and memory estimate for training a DFlash speculative decoding drafter on four NVIDIA RTX PRO 6000 Blackwell GPUs.

This message is deceptively simple: a single Python script that computes FLOPs, throughput, memory budgets, and training durations. But beneath the arithmetic lies a dense web of engineering judgment, architectural knowledge, and strategic decision-making. The assistant is not merely running numbers — it is validating the entire training plan against the unforgiving constraints of real hardware, answering the question that determines whether the project proceeds or pivots.

The Message

The assistant writes a single Python script that performs a comprehensive time and memory estimate for the DFlash training run. Here is the message in full:

# ============================================
# Time estimate for DFlash training on 4x PRO 6000
# ============================================

# Dataset stats
total_samples = 902_087
total_tokens = 1_866_000_000
mean_seq_len = 2068
loss_tokens = 1_633_000_000

# Hardware
num_gpus = 4
dp_pairs = 2  # 2 target + 2 drafter
bf16_tflops = 850  # PRO 6000 Blackwell BF16 dense
fp8_tflops = 1700  # PRO 6000 FP8 (if used for target)

# Model sizes  
target_params = 27e9
drafter_trainable = 1.7e9
drafter_frozen = 5e9  # embed_tokens + lm_head + verifier_lm_head + verifier_norm

# Training config
token_budget = 8192
block_size = 16
max_anchors = 512
epochs = 6

The script proceeds to compute samples per packed sequence, steps per epoch, FLOPs for the target forward pass (including the quadratic attention cost for 16 full-attention layers), drafter forward and backward times, per-step times for both BF16 and FP8 target precision, throughput in samples per second and tokens per second, time per epoch, total training time, and memory estimates for both the target and drafter GPUs.

The output reveals:

| Metric | BF16 Target | FP8 Target | |---|---|---| | Target forward time | 0.54s | 0.27s | | Drafter forward+backward | 98ms | 98ms | | Per-step time | 0.68s | 0.42s | | Throughput | 23,944 tok/s | 39,358 tok/s | | Time per epoch | 21.6h | ~13.2h | | Total 6 epochs | ~5.4 days | ~3.3 days |

The memory estimates show comfortable headroom: the target GPU uses ~64 GB out of 96 GB, while the drafter GPU uses only ~40 GB out of 96 GB.

Why This Message Was Written: The Strategic Context

To understand why the assistant wrote this message, we must look at the conversation leading up to it. In the preceding messages, the assistant had been deep in the weeds of debugging the DFlash training scripts. It had identified six critical bugs in the training code ([msg 7753]), including a drafter config that was incorrectly copying attention dimensions from the verifier model instead of using the correct Qwen3-style architecture. The assistant had also cross-referenced the z-lab reference implementation and discovered that the training script was missing sequence packing, noise augmentation, proper per-document anchor masking, correct position IDs, and torch.compile optimization.

But fixing bugs is only half the battle. Before provisioning expensive GPU instances and launching a multi-day training run, the assistant needed to answer a fundamental question: is this training run even feasible on the available hardware? The user had already been asked to choose between BF16 and FP8 target precision, and had selected BF16 (simpler, ~5.4 days). Now the assistant needed to validate that choice with hard numbers.

This message is the feasibility gate. It transforms vague concerns about "will it fit?" and "how long will it take?" into concrete, quantifiable answers. It is the moment where the project transitions from "we think this should work" to "we know this will work, and here is exactly how long it will take."

The Thinking Process: A Masterclass in Engineering Estimation

The assistant's reasoning, visible in the structure and content of the script, reveals a sophisticated mental model of the training pipeline. Let us unpack the layers of thinking embedded in this seemingly straightforward calculation.

Layer 1: Dataset-Aware Throughput Estimation

The assistant does not simply compute raw FLOPs and divide by peak throughput. It begins with the dataset statistics: 902,087 samples, 1.866 billion tokens, a mean sequence length of 2,068 tokens, and 1.633 billion loss tokens. These numbers are not decorative — they drive every downstream calculation.

The mean sequence length determines the average number of samples that can be packed into a single training sequence given the token budget of 8,192. The assistant computes avg_samples_per_pack = token_budget / mean_seq_len = 8192 / 2068 ≈ 4.0. This is a critical number: it means each packed sequence contains, on average, four documents. This directly multiplies the effective throughput — processing one packed sequence is equivalent to processing four individual samples, but with a single forward pass through the drafter.

The assistant also distinguishes between total tokens and loss tokens. The 233 million tokens excluded from the loss (1.866B - 1.633B) represent prompt tokens that are not used for training. This distinction matters for the loss computation but not for the FLOP count, since the forward pass processes all tokens regardless.

Layer 2: Hardware-Aware FLOP Accounting

The assistant's FLOP model is surprisingly detailed. For the target model forward pass, it computes:

target_flops = 2 * target_params * token_budget  # ~442 TFLOPS

This is the standard formula for a transformer forward pass: approximately 2 parameters tokens FLOPs (one multiply-add per parameter per token). But the assistant does not stop there. It adds the quadratic attention cost:

attn_flops_per_layer = 2 * 24 * token_budget * token_budget * 256
attn_flops = 16 * attn_flops_per_layer

This accounts for the 16 full-attention layers in the Qwen3.6-27B model (the remaining layers use a different attention mechanism). The formula 2 * heads * N * N * d captures the scaled dot-product attention computation: QK^T multiplication and PV multiplication, each contributing roughly heads * N * N * d FLOPs.

The total comes to 455.6 TFLOPS, which is only slightly more than the naive 442 TFLOPS estimate. This tells us that attention is not the dominant cost for this model at this sequence length — the linear layers dominate.

For the drafter, the assistant uses a simplified model focused on MLP FLOPs, noting that "flex attention sparsity" reduces the effective KV length. The drafter forward+backward is estimated at just 98ms — an order of magnitude faster than the target forward. This asymmetry is expected: the drafter has only 1.7B trainable parameters versus the target's 27B.

Layer 3: Throughput and Wall-Time Projection

The assistant then chains these per-step times into epoch-level and total training time estimates. The calculation is straightforward arithmetic, but the assumptions embedded in it are worth examining:

Layer 4: Memory Budgeting

The memory estimates reveal a comfortable margin on both GPU types:

Assumptions Made

Every estimate rests on assumptions, and this message is no exception. The assistant makes several notable assumptions that deserve scrutiny:

1. Peak TFLOPS as a Proxy for Real-World Performance

The assistant assumes the GPUs can sustain 850 BF16 TFLOPS (dense) and 1,700 FP8 TFLOPS. These are peak theoretical numbers. Real-world throughput depends on kernel efficiency, memory bandwidth, and operation fusion. The assistant's FLOP counts (455.6 TFLOPS for the target forward) are below the peak, suggesting the compute is not the bottleneck — but this assumes the kernels are well-optimized. If the target model uses custom GDN layers with suboptimal CUDA kernels, the actual throughput could be significantly lower.

2. Drafter Forward+Backward at 98ms

The drafter time estimate is notably rough. The assistant computes only MLP FLOPs and ignores attention entirely, with a note that "flex attention sparsity" reduces the effective KV length. The 98ms estimate is plausible but unvalidated — it assumes the drafter's attention is efficiently handled by flex attention kernels, which may not be true on Blackwell hardware (sm_120) with early versions of PyTorch or Triton.

3. Linear Scaling with Data Parallelism

The assistant assumes perfect scaling with 2 DP pairs: each pair processes half the data, and there is no communication overhead. In practice, gradient synchronization across DP pairs adds some overhead, and the assistant's 50ms overhead estimate may not fully capture this.

4. Packing Efficiency

The assistant assumes that packing 4 samples per sequence is always possible and that the packing does not introduce computational overhead. In practice, packing requires careful padding and masking, and the effective throughput may be lower if the packing algorithm cannot achieve the theoretical 4x density.

5. Stable Training Dynamics

The estimate assumes the training runs to completion without interruptions, crashes, or convergence issues. In reality, training runs often require restarts, hyperparameter tuning, and debugging — the 5.4-day estimate is best-case wall time, not expected wall time.

Input Knowledge Required

To understand this message fully, a reader needs knowledge spanning several domains:

ML Hardware Architecture

Transformer Model Architecture

DFlash Training Pipeline

Engineering Estimation

Output Knowledge Created

This message produces several concrete outputs that shape the subsequent conversation:

1. Feasibility Confirmation

The estimate confirms that the training run is feasible on 4× RTX PRO 6000 GPUs. Both the compute and memory budgets are within bounds, with comfortable headroom on the drafter GPU and adequate headroom on the target GPU.

2. Timeline Projection

The 5.4-day (BF16) estimate gives the team a concrete timeline for planning. This is long enough to warrant careful monitoring and checkpointing but short enough to be practical for a research project.

3. Bottleneck Identification

The estimate reveals that the target model forward pass dominates the per-step time (0.54s out of 0.68s for BF16). The drafter is not the bottleneck. This informs optimization priorities: if the training needs to be faster, the target model forward pass is where to focus.

4. Resource Allocation Validation

The memory estimates validate the 2-target + 2-drafter GPU allocation. The target GPUs are well-utilized (~64 GB), while the drafter GPUs have significant headroom (~40 GB). This confirms that the allocation is reasonable and that the drafter GPUs could potentially handle larger batch sizes or additional model components.

5. Decision Support for Precision Choice

The estimate quantifies the FP8 advantage (~40% speedup, from 5.4 to 3.3 days) but also shows that BF16 is viable. Since the user chose BF16, this estimate provides the rationale for that choice: the additional complexity of FP8 quantization is not worth the time savings when BF16 already fits comfortably.

Decisions Made (and Not Made)

This message is primarily analytical rather than decisional. It does not change the training plan or introduce new directions. However, it does serve as the foundation for several implicit decisions:

Implicit Decision: Proceed with BF16

The estimate confirms that BF16 training is feasible within a reasonable timeframe. Combined with the user's explicit choice of BF16, this message validates that decision with hard numbers.

Implicit Decision: No Need for FP8 Optimization

The comfortable memory headroom and acceptable training time mean there is no pressure to pursue FP8 quantization. The team can proceed with the simpler BF16 path.

Implicit Decision: Acceptable Timeline

The 5.4-day estimate is presented without commentary on whether this is acceptable. By not flagging it as problematic, the assistant implicitly signals that this timeline is reasonable for the project's goals.

Decision Not Made: Whether to Reduce Epochs

The assistant considers this possibility in the subsequent message ([msg 7755]: "Let me also think about whether reducing the epoch count could help, since 6 epochs at 5.4 days is substantial") but does not act on it here. This message simply presents the numbers; the decision to adjust hyperparameters comes later.

Mistakes and Incorrect Assumptions

While the estimate is generally sound, several assumptions deserve critical examination:

1. Optimistic TFLOPS Utilization

The assistant assumes the target model can sustain 850 BF16 TFLOPS, which is the theoretical peak for the RTX PRO 6000 Blackwell. Real-world utilization for transformer inference typically ranges from 30-60% of peak, depending on kernel efficiency and memory bandwidth. If the actual utilization is 50%, the target forward time would double to ~1.08s, and the total training time would increase to ~7.5 days.

2. Overhead Underestimation

The 50ms overhead estimate is likely too low for a complex training pipeline involving data loading, packing, masking, and inter-GPU communication. Real-world overhead for such pipelines often reaches 100-200ms per step.

3. Drafter Time Approximation

The drafter time estimate ignores attention FLOPs entirely. While flex attention reduces the computational cost, it does not eliminate it. The 98ms estimate may be optimistic by 20-50%.

4. No Accounting for Checkpointing

The estimate does not include time for saving checkpoints (every 5,000 steps), which involves writing ~13 GB of drafter weights to disk. At typical NVMe write speeds, this adds ~5-10 seconds per checkpoint, or ~20-40 minutes over the full training run.

5. Ignoring Data Pipeline Bottlenecks

The estimate assumes the data pipeline can keep up with the GPU throughput. With 19 GB of tokenized data and a 9-minute sync time (as noted in the chunk summary), the data loading could become a bottleneck if not properly cached in memory.

The Broader Significance

This message exemplifies a critical skill in large-scale ML engineering: the ability to translate a complex training pipeline into a simple, verifiable computational model. The assistant does not need to run the training to know it will work — it can reason about FLOPs, memory, and throughput with sufficient accuracy to make confident decisions.

This is the difference between a junior engineer who launches a training job and hopes it works, and a senior engineer who knows it will work because they have done the math. The assistant's estimate is not perfect — it makes simplifying assumptions and ignores second-order effects — but it is good enough. It answers the essential questions: will it fit? how long will it take? where are the bottlenecks?

In the subsequent conversation ([msg 7755]), the assistant uses this estimate to present a comprehensive training plan with phases, timelines, and contingency strategies. The estimate becomes the backbone of the entire execution plan. When the training eventually launches and runs for days, it is this message — this single Python script — that gave the team the confidence to proceed.

Conclusion

Message [msg 7754] is a masterclass in ML engineering estimation. It demonstrates that rigorous back-of-the-envelope calculations are not a substitute for running code, but they are an essential prerequisite. Before provisioning GPUs, before downloading models, before debugging Triton autotuner crashes, the assistant asks: does the math work?

The answer, computed in a single Python script, is a confident yes. The training will take approximately 5.4 days on 4× RTX PRO 6000 GPUs with BF16 precision. The memory fits with comfortable headroom. The throughput is adequate. The project can proceed.

This message is the quiet hero of the session — the analytical foundation upon which the entire training run was built. It is a reminder that in the world of large-scale ML, the most important tool is not a faster GPU or a better optimizer, but a clear understanding of what the hardware can do and what the software requires. The rest is just execution.