perf(fix): accumulate per-microbatch FLOPS metadata for accurate… - #3529
Conversation
… VLM TFLOPS - Replace static cfg.model.seq_length with accumulated actual padded lengths from each micro-batch (_flops_seqlen_sum/sq_sum). - Add ViT encoder FLOPS via _flops_vision_patches from visual grid THW. - Scale per-rank accumulators by dp_size for global FLOPS estimation. - Revert train_step return signature; zero API surface changes. - Also support trace_prefix in pytorch profiler initialization. Signed-off-by: dingtianwei.dtw <dingtianwei.dtw@alibaba-inc.com>
0a5474e to
7f1ff31
Compare
yaoyu-33
added
area:perf
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
ReviewBug —
Missing tests The PR description says "Added |
Bug:
|
yaoyu-33
added
waiting-on-customer
1. Wire seqlen_squared_sum into core-attn via core_attn_seq_factor (was dead code). 2. Refactor vit_flops to take (cfg, batch_size, num_patches) so ViT hyperparameters are read from cfg.model.vision_config. 3. Align training_log TFLOP/s/GPU with the main-loop FLOPS accumulator for variable-length batches. 4. Tighten seqlen_sum typing to int | None and add some unit tests. Signed-off-by: dingtianwei.dtw <dingtianwei.dtw@alibaba-inc.com>
…uared_sum" . Signed-off-by: dingtianwei.dtw <dingtianwei.dtw@alibaba-inc.com>
svcnvidia-nemo-ci
removed
the
waiting-on-customer
84c7535 to
f5d422b
Compare
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
|
/ok to test f5d422b |
svcnvidia-nemo-ci
removed
the
waiting-on-maintainers
yaoyu-33
added
the
ready-to-merge
|
@SophusDavid : please run precommit locally to resolve lint issues |
svcnvidia-nemo-ci
added
waiting-on-customer
Apologies — I didn't notice the UT job was failing on this PR earlier. I've now run the relevant tests locally and pushed fixes:
While investigating, I also found that this patch was computing seqlen_sum / seqlen_squared_sum / num_vision_patches incorrectly under VPP (interleaved pipeline): forward_step_func is invoked once per virtual-stage per microbatch, so the accumulators were inflated by vp_size, while num_floating_point_operations() already covers all layers of the full model. This caused TFLOP/s/GPU to be reported vp_size× too high (e.g. 3× with vp_size=3). |
svcnvidia-nemo-ci
removed
the
waiting-on-customer
|
/ok to test b7cf908 |
|
/ok to test 0369c27 |
| """ | ||
| if num_vision_patches <= 0: | ||
| return 0 | ||
| patches_per_image = num_vision_patches / batch_size if batch_size > 0 else num_vision_patches |
There was a problem hiding this comment.
this has an implicit assumption that batch size == num images?
cuichenx
removed
the
ready-to-merge
aroshanghias-nvd
added
the
waiting-on-customer
|
merging this first, comment above can be fixed in a followup PR |
svcnvidia-nemo-ci
removed
the
waiting-on-customer
Packed THD training (offline-packed LLM SFT and VLM in-batch packing) over-counts attention FLOPS by treating the whole pack as one length- seq_length sequence (pack_length²). Actual attention work is Σᵢ sᵢ² over the real sub-sequence lengths. New helper accumulate_flops_metadata() in flop_utils.py extracts the real sub-seq lengths from cu_seqlens (preferring cu_seqlens_unpadded when present) and feeds Σᵢ sᵢ² into the existing seqlen_squared_sum accumulator from NVIDIA-NeMo#3529. Falls back to BSHD mbs * seq_len² when no cu_seqlens is provided — bit-exact identical to legacy on dense pretraining and non-packed paths. Wired into gpt_step, vlm_step, qwen3_vl_step, and qwen3_omni_step. Verified on cw-dfw (same seed, same data, same iter times, identical loss values across paired runs — only the reported TFLOPS differs): - qwen3_8b_sft seq=2048: baseline 162.6 vs fix 155.8 TFLOP/s/GPU (+4%) - qwen3_8b_sft seq=4096: baseline 339.9 vs fix 156.7 TFLOP/s/GPU (+117%) - qwen35_vl_9b_sft : baseline 261.6 vs fix 88.9 TFLOP/s/GPU (+194%) The seq=2048→4096 pair on the same LLM recipe is the cleanest demonstration: the fix is near-flat (155.8 vs 156.7 — attention work is determined by per-sample lengths, not pack length) while the baseline doubles because its pack_length² scales quadratically. 9 new unit tests in test_flop_utils.py::TestAccumulateFlopsMetadata cover the BSHD fallback, THD with cu_seqlens, padded cu_seqlens via cu_seqlens_argmin, cu_seqlens_unpadded precedence, additive accumulation, and the regression headline (32-sample pack → 32x smaller attention work than BSHD approximation). Signed-off-by: Chen Cui <chcui@nvidia.com>
Packed THD training (offline-packed LLM SFT and VLM in-batch packing) over-counts attention FLOPS by treating the whole pack as one length- seq_length sequence (pack_length²). Actual attention work is Σᵢ sᵢ² over the real sub-sequence lengths. New helper accumulate_flops_metadata() in flop_utils.py extracts the real sub-seq lengths from cu_seqlens (preferring cu_seqlens_unpadded when present) and feeds Σᵢ sᵢ² into the existing seqlen_squared_sum accumulator from NVIDIA-NeMo#3529. Falls back to BSHD mbs * seq_len² when no cu_seqlens is provided — bit-exact identical to legacy on dense pretraining and non-packed paths. Wired into gpt_step, vlm_step, qwen3_vl_step, and qwen3_omni_step. Verified on cw-dfw (same seed, same data, same iter times, identical loss values across paired runs — only the reported TFLOPS differs): - qwen3_8b_sft seq=2048: baseline 162.6 vs fix 155.8 TFLOP/s/GPU (+4%) - qwen3_8b_sft seq=4096: baseline 339.9 vs fix 156.7 TFLOP/s/GPU (+117%) - qwen35_vl_9b_sft : baseline 261.6 vs fix 88.9 TFLOP/s/GPU (+194%) The seq=2048→4096 pair on the same LLM recipe is the cleanest demonstration: the fix is near-flat (155.8 vs 156.7 — attention work is determined by per-sample lengths, not pack length) while the baseline doubles because its pack_length² scales quadratically. 9 new unit tests in test_flop_utils.py::TestAccumulateFlopsMetadata cover the BSHD fallback, THD with cu_seqlens, padded cu_seqlens via cu_seqlens_argmin, cu_seqlens_unpadded precedence, additive accumulation, and the regression headline (32-sample pack → 32x smaller attention work than BSHD approximation). Signed-off-by: Chen Cui <chcui@nvidia.com>
Packed THD training (offline-packed LLM SFT and VLM in-batch packing) over-counts attention FLOPS by treating the whole pack as one length- seq_length sequence (pack_length²). Actual attention work is Σᵢ sᵢ² over the real sub-sequence lengths. New helper accumulate_flops_metadata() in flop_utils.py extracts the real sub-seq lengths from cu_seqlens (preferring cu_seqlens_unpadded when present) and feeds Σᵢ sᵢ² into the existing seqlen_squared_sum accumulator from NVIDIA-NeMo#3529. Falls back to BSHD mbs * seq_len² when no cu_seqlens is provided — bit-exact identical to legacy on dense pretraining and non-packed paths. Wired into gpt_step, vlm_step, qwen3_vl_step, and qwen3_omni_step. Verified on cw-dfw (same seed, same data, same iter times, identical loss values across paired runs — only the reported TFLOPS differs): - qwen3_8b_sft seq=2048: baseline 162.6 vs fix 155.8 TFLOP/s/GPU (+4%) - qwen3_8b_sft seq=4096: baseline 339.9 vs fix 156.7 TFLOP/s/GPU (+117%) - qwen35_vl_9b_sft : baseline 261.6 vs fix 88.9 TFLOP/s/GPU (+194%) The seq=2048→4096 pair on the same LLM recipe is the cleanest demonstration: the fix is near-flat (155.8 vs 156.7 — attention work is determined by per-sample lengths, not pack length) while the baseline doubles because its pack_length² scales quadratically. 9 new unit tests in test_flop_utils.py::TestAccumulateFlopsMetadata cover the BSHD fallback, THD with cu_seqlens, padded cu_seqlens via cu_seqlens_argmin, cu_seqlens_unpadded precedence, additive accumulation, and the regression headline (32-sample pack → 32x smaller attention work than BSHD approximation). Signed-off-by: Chen Cui <chcui@nvidia.com>
…IDIA-NeMo#3529) Signed-off-by: dingtianwei.dtw <dingtianwei.dtw@alibaba-inc.com> Signed-off-by: meatybobby <meatybobby@gmail.com> Co-authored-by: meatybobby <meatybobby@gmail.com> Co-authored-by: aroshanghias-nvd <aroshanghias@nvidia.com> Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
7 participants
perf(training): improve VLM TFLOPS calculation with actual sequence length and ViT
What does this PR do ?
Fixes inaccurate TFLOPS reporting for VLM training where the reported FLOPS was based on the statically configured
cfg.model.seq_lengthinstead of the dynamically padded sequence lengths that FlashAttention actually processes. This led to significant overestimation (up to ~2.6x for typical VLM SFT scenarios with short sequences + images).The PR introduces a per-micro-batch accumulator pattern that captures the real padded sequence length and vision patch count from each
forward_stepinvocation, scales them by data-parallel size for global estimation, and feeds them into an extendednum_floating_point_operations()API.Closes #3498
Changelog
vlm_step.py
_flops_seqlen_sumand_flops_seqlen_sq_sumacross micro-batches usingtokens.shape[1](actual padded length)._flops_vision_patchesfromvisual_inputs.image_grid_thw/video_grid_thw."total_tokens"topacked_seq_paramsfor downstream visibility.train.py
_flops_*accumulators to 0 before eachwrapped_train_stepcall.dp_sizefor global FLOPS, and pass effectiveseq_lengthtotraining_log.seq_lengthreturn value fromtrain_step(reverting to original 8-element return signature).flop_utils.py
vit_flops()helper for Vision Transformer encoder FLOPS (bidirectional attention + GELU MLP + patch merger).num_floating_point_operations()signature with optionalseqlen_sum,seqlen_squared_sum, andnum_vision_patchesparameters._compute_vit_flops()that derives per-image patch count from total batch patches for correct quadratic attention scaling.train_utils.py
training_log()now acceptsseq_lengthparameter and uses accumulated_flops_vision_patches(scaled bydata_parallel_size) for consistent per-log FLOPS recomputation.Before your PR is "Ready for review"
Pre checks:
test_flops_calculation.pywith 13 verification tests covering VLM and non-VLM scenarios.Additional Information