Revert "Fix LLAMA3 LoRa TFLOPs Formula (#2416)"

[ko3n1g]

This reverts commit fa4a01b.

What does this PR do ?

Failing CI https://github.com/NVIDIA-NeMo/Megatron-Bridge/actions/runs/22419432650/job/64916324689?pr=2535#step:4:5389

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:

• 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

• Related to # (issue)

Summary by CodeRabbit

• Refactor
  • Streamlined floating-point operations calculation by removing specialized branching logic and simplifying pathways to invoke custom model-provided calculation methods.

[coderabbitai]

Caution

Review failed

The pull request is closed.

ℹ️ Recent 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 3919026 and 0c5a123.

📒 Files selected for processing (1)

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

---

📝 Walkthrough

Walkthrough

This PR removes LoRA-specific branching logic from the FLOPs calculation utility by simplifying the num_floating_point_operations function to consistently delegate to a model-provided _get_num_floating_point_operations method when available, eliminating conditional LoRA discrimination checks and related LoRA-dependent stats computation paths.

Changes

Cohort / File(s)	Summary
FLOPs Calculation Simplification src/megatron/bridge/training/utils/flop_utils.py	Removed LoRA-specific imports, type checks, and conditional branching in num_floating_point_operations and transformer_flops. The function now directly uses the model's custom FLOPs calculation method, eliminating LoRA discrimination logic and associated seq_length validation paths.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

Possibly related PRs

• PR #2416: Inverse changes that add LoRA-aware FLOPs branching logic and skip the model method when LoRA is active, directly modifying the same flop_utils transformer FLOPs logic.

Suggested labels

r0.3.0

Suggested reviewers

• guyueh1
• erhoo82

✨ Finishing Touches

• 📝 Generate docstrings (stacked PR)
• 📝 Generate docstrings (commit on current branch)

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch ko3n1g/revert/fa4a01b798a1e3ab4f08367bba68244976fe9a7d

---

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.}