Skip to content

[draft] feat: add flextron support to Hybrid models#2474

Draft
rohitrango wants to merge 15 commits into
NVIDIA-NeMo:mainfrom
rohitrango:rohit/flextron
Draft

[draft] feat: add flextron support to Hybrid models#2474
rohitrango wants to merge 15 commits into
NVIDIA-NeMo:mainfrom
rohitrango:rohit/flextron

Conversation

@rohitrango
Copy link
Copy Markdown
Contributor

What does this PR do ?

Adds flextron support for hybrid mamba models (mcore generation backend)

Usage

uv run examples/grpo.py --config examples/configs/grpo_math_1B_flex.yaml

# Add a code snippet demonstrating how to use this

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 run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
  • Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

rohitrango added 5 commits May 8, 2026 15:40
Signed-off-by: rohitrango <rohit.rango@gmail.com>
Signed-off-by: rohitrango <rohit.rango@gmail.com>
Signed-off-by: rohitrango <rohit.rango@gmail.com>
Signed-off-by: rohitrango <rohit.rango@gmail.com>
Signed-off-by: rohitrango <rohit.rango@gmail.com>
@copy-pr-bot
Copy link
Copy Markdown

copy-pr-bot Bot commented May 12, 2026

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.

rohitrango added 10 commits May 12, 2026 09:36
Signed-off-by: rohitrango <rohit.rango@gmail.com>
Signed-off-by: rohitrango <rohit.rango@gmail.com>
Signed-off-by: rohitrango <rohit.rango@gmail.com>
- fix no flextron bug
- commented out disable forward hook logic
- 8 node config

Signed-off-by: rohitrango <rohit.rango@gmail.com>
Signed-off-by: rohitrango <rohit.rango@gmail.com>
Signed-off-by: rohitrango <rohit.rango@gmail.com>
Signed-off-by: rohitrango <rohit.rango@gmail.com>
Signed-off-by: rohitrango <rohit.rango@gmail.com>
Signed-off-by: rohitrango <rohit.rango@gmail.com>
Signed-off-by: rohitrango <rohit.rango@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant