Add vision flops calculation for vlm

[JeffPengCoder]

What does this PR do ?

As there is a trend that vision part in vlm is becoming larger and larger, we could not neglect flops costs on vision part any more. Especially MoE language model is becoming more popular, compared with activated weights, we could speculate that compute in vision part would become more understanding.

current flops calculation does not include vision part for vlm, this PR fill the gap

Changelog

• Add specific line by line info of high level changes in this PR.

GitHub Actions CI

See the CI sectionin the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.

Before your PR is "Ready for review"

Pre checks:

• [Y] Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

There are several options to do this functionality.

1. calculate flops upon vision forward completes;
2. record neccessary info and do calculation at the end of step, this PR follow this method

• Related to # (issue)

Summary by CodeRabbit

• New Features

  • Added vision tower floating-point operation (FLOP) calculation for accurate performance metrics in vision-language models.
  • Introduced runtime metric collection for per-sample language and vision token statistics to enhance monitoring.

• Bug Fixes

  • Improved configuration validation to robustly handle diverse configuration formats and edge cases in finalization flows.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

📝 Walkthrough

Walkthrough

The PR adds vision FLOPs instrumentation to training pipelines and improves configuration type-safety. Runtime statistics for vision-language models are captured during forward passes, configuration objects receive safe type normalization and finalization guards, and FLOP calculation utilities now estimate vision tower costs using both runtime-derived and config-based assumptions.

Changes

Cohort / File(s)	Summary
Configuration Type-Safety src/megatron/bridge/training/config.py	Added guards to only finalize profiling, nvrx_straggler, and tensor_inspect if they have a finalize method; converts mixed_precision and comm_overlap from dicts to their respective Config objects before finalization in runtime_config_update().
Vision FLOP Calculation src/megatron/bridge/training/utils/flop_utils.py, src/megatron/bridge/training/vlm_step.py	Introduces vision_floating_point_operations() function to estimate FLOPs for vision components (patch embedding, ViT blocks, patch mergers). Integrates runtime statistics collection in forward_step to populate per-sample vision token counts and squared length sums. Calculates vision FLOPs using runtime metrics when available, falling back to config-based assumptions; multiplies by 3.0 to account for backward pass cost.

Sequence Diagram

sequenceDiagram
    participant FS as forward_step
    participant CFG as Config Object
    participant FLOP as FLOP Utils
    participant OUT as Output FLOPs
    
    FS->>FS: Process vision/language inputs
    FS->>CFG: Store runtime statistics<br/>(_runtime_vision_tokens_pre_per_sample, etc.)
    FS->>FS: Continue training forward pass
    
    FLOP->>CFG: Check for vision_config
    FLOP->>CFG: Retrieve runtime statistics if available
    FLOP->>FLOP: Calculate vision FLOPs<br/>(patch embedding, ViT, mergers)
    FLOP->>FLOP: Apply training multiplier (3.0×)
    FLOP->>FLOP: Integrate with hybrid/standard<br/>transformer FLOPs
    FLOP->>OUT: Return total FLOPs

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

• Bridging perf from NeMo2 to Mbridge for certain configs #2199: Modifies TrainingConfig validation alongside this PR's changes to validate() and runtime_config_update() methods.

Suggested labels

r0.3.0

Suggested reviewers

• ananthsub

🚥 Pre-merge checks | ✅ 3 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Test Results For Major Changes	⚠️ Warning	PR adds major vision FLOPs calculation feature (+252 net lines) but lacks test results, performance validation, or numeric benchmarks despite PR objectives requiring testing updates.	Add unit tests for vision_floating_point_operations function, integration tests for vlm_step.py instrumentation, and document validation results in PR description.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'Add vision flops calculation for vlm' directly and specifically describes the main objective of the PR—adding FLOPs calculation for the vision component in vision-language models.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 3

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@src/megatron/bridge/training/config.py`:
- Around line 1910-1919: The current ordering causes
cfg.set_data_parallel_size() to access cfg.comm_overlap as if it were an object
while it may still be a dict; before calling cfg.set_data_parallel_size()
normalize cfg.comm_overlap: if cfg.comm_overlap is a dict, replace it with
CommOverlapConfig(**cfg.comm_overlap), and if it has finalize(), call finalize()
so that cfg.comm_overlap is a proper object; then call
cfg.set_data_parallel_size() and only after that call
cfg.comm_overlap.setup(cfg.model, cfg.optimizer, cfg.ddp) (keeping guards for
cfg.comm_overlap is not None).

In `@src/megatron/bridge/training/vlm_step.py`:
- Around line 459-488: The vision FLOPs/statistics computation using vision_cfg,
image_grid_thw, video_grid_thw, and microbatch_size_for_stats should be made
fail-open so telemetry errors never abort training: wrap the entire block that
computes vision_tokens_pre_total, vision_sum_seqlen_sq_pre_total,
vision_tokens_pre_per_sample, vision_tokens_post_per_sample and sets
state.cfg._runtime_vision_* in a try/except, catch Exception, log the exception,
and on any error set the runtime vision stats to safe defaults (e.g., 0) so
forward continues; ensure spatial_merge_size, hw/t, and prod(dim=-1) uses remain
unchanged inside the try and do not propagate exceptions out of the function.
- Around line 424-450: Replace the broad "except Exception" in the block that
inspects forward_args/packed (variables: packed, cu, forward_args, tokens) and
writes to state.cfg._runtime_lm_* with a targeted exception handler that only
catches the expected errors (AttributeError, IndexError, RuntimeError); keep the
try body limited to the attribute/tensor operations (accessing
packed.cu_seqlens_q_padded, indexing cu, computing lengths and tensors) and move
any pure-assignment work into an else branch where possible, and ensure the
fallback sets microbatch_size_for_stats the same way as currently done when
those specific exceptions occur.

---

ℹ️ Review info

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between a3e44b4 and 3d52466.

📒 Files selected for processing (3)

• src/megatron/bridge/training/config.py
• src/megatron/bridge/training/utils/flop_utils.py
• src/megatron/bridge/training/vlm_step.py

[JeffPengCoder]

Add one more comment

After adding vision flops, for qwen3-vl-30b-a3b, (vison 545m, activate 3b，545m : 3b ~ 1 :6), total flops increase from 494 to 554(60 : 494 ~ 1 : 8), as text length would be >0, increased flops would not exceed 1/6, current increase is as expected.

For qwen3-30b-a3b, flops almost the same between before code change and after code change.

[yaoyu-33]

Closing as superseded by #3529. Current main includes vision-side FLOPs accounting and per-microbatch metadata with focused tests.