gpt-oss Mcore support

[terrykong]

related to #841

[ashors1]

I'm working on pipe-cleaning GPT-OSS using Yu's megatron bridge PR and specific GPT-OSS megatron-bridge and megatron branches. Currently, I'm running into a few challenges:

• I need a newer version of VLLM to pull in GPT-OSS support, which also requires newer ray. Currently working through these dependency issues
• I tried running GPT-OSS with SFT to pipe clean the Megatron side and I hit OOM. This is probably related to the 70B OOM that @terrykong mentioned here

Will continue to provide updates as I have them

[kwanUm]

Thanks for tackling GPT-OSS support — this would be a game changer for tool-based RL.

I tried wiring GPT-OSS 20B with NeMo-RL + vLLM and got stuck at the refit path during weight streaming:

• Site: nemo_rl/models/generation/vllm_backend.py → VllmInternalWorkerExtension.update_weights_from_local_ipc_handles
• Error: Attempted to load weight (torch.Size([32, 5760])) into parameter (torch.Size([32, 5888]))
• Also seeing many incoming HF weights not found on the vLLM side (e.g., HF: model.embed_tokens, self_attn.{q_proj,k_proj,v_proj}, input/post_attention_layernorm, experts.{gate_up_proj,down_proj} vs vLLM: model.embedding, attn.qkv (fused), attn.norm/mlp.norm, mlp.experts.{w13_, w2_}, router.*). This looks like vLLM’s GPT-OSS implementation expects fused QKV and 128-aligned dims/padding, while the HF training model exposes unfused heads and unpadded shapes.

Is someone on the Nemo RL team prioritizing this?

[terrykong]

Hi @kwanUm . Thanks for sharing your use case, it's possible some work is needed on dtensor, could I ask you to continue this discussion over on #841 since that tracks the dtensor support which is a parallel effort?

[ashors1]

Quick update here: I was able to resolve the dependency and OOM issues we were seeing before. There is currently a bug in the Megatron Bridge mcore --> HF conversion that arises when EP > 1. I have a proposed fix which the Megatron Bridge team is helping push to the gpt-oss branch. Now, training wtih gpt-oss 20b looks okay. Rewards are increasing, but logprob error is pretty consistently > 2. I'm following up with some folks to understand what a reasonable logprob error is for this model, and will work on creating a more minimal reproducer for the high logprob error if required.

[euronymous-aithal]

@yaoyu-33 :it seems error in the logits b/w MCore and vLLM is more than the accepted threshold ? can you or Chen help on this?

[kzhou92]

Hi @ashors1 , thanks a lot for the development. Just wonder how is the progress on the development. Any possibility of sharing an early version?

[ashors1]

Hi @kzhou92, thanks for your interest. GPT-OSS support in Megatron has not been open-sourced yet, and the VLLM vs Mcore discrepancy still remains. I have been away this week, but I will be prioritizing investigating the VLLM/Mcore logprob issues in the coming days. We will keep this issue up-to-date on any progress.

[terrykong]

re-opening @yaoyu-33 since we still need to validate E2E + update submodule