Refit regression for MoE models using Megatron-Bridge

[yfw]

Refit is slower with Megatron-Bridge than nemo.tron for moe models. To reproduce, use these branches:

baseline: https://github.com/NVIDIA-NeMo/RL/tree/yifu/before_mbridge
mbridge: https://github.com/NVIDIA-NeMo/RL/tree/yifu/mbridge_dsv3

DSv2-Lite:

uv run python examples/run_grpo_math.py --config=examples/configs/grpo_math_1B_megatron.yaml grpo.val_batch_size=2 policy.model_name=deepseek-ai/DeepSeek-V2-Lite-Chat cluster.gpus_per_node=8 policy.megatron_cfg.pipeline_model_parallel_size=4 policy.megatron_cfg.num_layers_in_first_pipeline_stage=7 policy.megatron_cfg.num_layers_in_last_pipeline_stage=6 policy.max_total_sequence_length=1024 checkpointing.enabled=False checkpointing.save_period=5 grpo.val_period=-1 grpo.val_at_start=False grpo.max_val_samples=16 policy.megatron_cfg.expert_model_parallel_size=2 policy.megatron_cfg.apply_rope_fusion=False

Make sure to set a different NRL_MEGATRON_CHECKPOINT_DIR when testing with mbridge so it won't reuse the nemo.tron checkpoints.

[terrykong]

@ZhiyuLi-Nvidia is chasing a memory leak regression in mbridge preventing us from using the optimal number of refit buckets

[guyueh1]

can we close this? are we waiting for any PRs? @ZhiyuLi-Nvidia

[ZhiyuLi-Nvidia]

can we close this? are we waiting for any PRs? @ZhiyuLi-Nvidia

We need some PR
@guyueh1 do you plan to use your change?
guyueh1/Megatron-LM@af73aa2

Otherwise I plan to disable cache here:
There's LRU cache in YarnRope (Mcore), which avoid garbage collection after offloading.

-    @lru_cache(maxsize=32)
     def forward(self, max_seq_len: int, offset: int = 0, packed_seq: bool = False) -> Tensor:
         """Forward pass of Yarn Rotary Embedding.

[guyueh1]

@ZhiyuLi-Nvidia I linked the PR to move lru_cache to a function instead of a method, I believe it is targeting v0.4
what about the other fix in m-bridge side, can you copy the link here?

[ZhiyuLi-Nvidia]

@ZhiyuLi-Nvidia I linked the PR to move lru_cache to a function instead of a method, I believe it is targeting v0.4
what about the other fix in m-bridge side, can you copy the link here?

Yeap, here's the cache optimization PR: #1187