Revert "[training, perf] fix: THD-aware FLOPS via cu_seqlens (Σᵢ sᵢ²) (re-land of #3839)"

[malay-nagda]

Reverts #4366

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[ko3n1g]

FLOP calc now seems to be broken for MoE models too.

[claude]

Bug: order-of-operations mismatch with train.py.

Here num_vision_patches is scaled by dp_size before the VPP correction (line 1108 divides the already-scaled value). In train.py the order is reversed — VPP correction first, then * dp_size.

With integer division these produce different results. E.g. local=3, vp=2, dp=4:

• train.py: (3 // 2) * 4 = 4
• here: (3 * 4) // 2 = 6

Should match train.py's order: VPP-correct the local value first, scale by dp_size after.

Suggested change

	num_vision_patches = local_vision_patches * config.data_parallel_size if local_vision_patches > 0 else 0
	vp_size = getattr(config.model, "virtual_pipeline_model_parallel_size", None)
	if isinstance(vp_size, int) and vp_size > 1:
	local_seqlen_sum = local_seqlen_sum // vp_size
	local_seqlen_sq_sum = local_seqlen_sq_sum // vp_size
	num_vision_patches = num_vision_patches // vp_size
	num_vision_patches = local_vision_patches
	vp_size = getattr(config.model, "virtual_pipeline_model_parallel_size", None)
	if isinstance(vp_size, int) and vp_size > 1:
	local_seqlen_sum = local_seqlen_sum // vp_size
	local_seqlen_sq_sum = local_seqlen_sq_sum // vp_size
	num_vision_patches = num_vision_patches // vp_size

[claude]

And the dp_size scaling for num_vision_patches should happen after the VPP block (to match train.py):

Suggested change

	if local_seqlen_sum > 0:
	seqlen_sum = local_seqlen_sum * config.data_parallel_size
	seqlen_squared_sum = local_seqlen_sq_sum * config.data_parallel_size
	if local_seqlen_sum > 0:
	seqlen_sum = local_seqlen_sum * config.data_parallel_size
	seqlen_squared_sum = local_seqlen_sq_sum * config.data_parallel_size
	num_vision_patches = num_vision_patches * config.data_parallel_size if num_vision_patches > 0 else 0

[claude]

test

[claude]

Review - Revert of #4366 (THD-aware FLOPS)

[claude]

Review - Revert of #4366 (THD-aware FLOPS)

Bug

train_utils.py / train.py vision-patch VPP ordering mismatch - In training_log (train_utils.py), num_vision_patches is scaled by dp_size before the VPP // vp_size correction (lines 1102-1108). In train.py the order is reversed (VPP first, then * dp_size). Integer division makes these give different results when vp_size > 1 and vision patches are present. See inline comment for a suggested fix.

Observations (non-blocking)

• Substantial code duplication between train.py:549-584 and train_utils.py:1093-1131: accumulator reading, MagicMock coercion, VPP correction, and dp_size scaling are repeated almost verbatim. Consider extracting a lightweight shared function.

• vlm_step.py:498 calls .item() per micro-batch, introducing a GPU-to-CPU sync on every micro-batch with vision inputs. The removed accumulate_flops_metadata kept tensors on-device to avoid this.

• training_log throughput now reflects only the last step when log_interval > 1 (accumulators are reset each step in train.py:440-442). The removed interval-delta approach averaged over the full logging window. Under variable-length batches the single-step number can be noisy.

• test_throughput_uses_interval_flops_delta was removed but no replacement test covers the new training_log FLOPS path.

Suggested test cases

No perf tests impacted.