[megatron-policy] refit fails on Qwen3.5-MoE A17B with PP > 1: "Object present on multiple PP ranks" due to MTP shared embedding
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
### TL;DR
Running NeMo-RL GRPO on `Qwen/Qwen3.5-397B-A17B` (Megatron PP=8 / EP=32, vLLM
colocated, 32 nodes) on top of the recipe shipped at
`examples/configs/recipes/llm/grpo-qwen3.5-397ba17b-32n8g-megatron.yaml`
fails at setup with `ValueError: Object present on multiple PP ranks: [0, 7]`.
This is the **first** of three failures that have to be cleared before the
recipe runs end-to-end. The journey to a passing 10-step smoke test (with
all three issues worked around) is logged here in chronological order so
maintainers can decide which of them deserve an upstream fix, which are
container-version bumps, and which are documentation gaps.
Final verified config that completed 9/10 GRPO steps cleanly
(Step 10 cut by SLURM 4h time limit, mid-training, not by any of the issues
below):
| Knob | Recipe default | What worked |
|---|---|---|
| `policy.generation.vllm_cfg.expert_parallel_size` (`VLLM_EP`) | 64 | **256** |
| `policy.generation.vllm_cfg.enforce_eager` | False | **True** |
| `NRL_REFIT_BUFFER_MEMORY_RATIO` (env) | 0.3 (default) | **0.39** |
| `vllm_worker.py` sleep level | `level=1` | **`level=2`** (bind-mount patch) |
| `megatron_policy_worker.py::_calculate_refit_param_info` | unpatched | drops MTP-replicated `embedding.word_embeddings.weight` on last PP stage (bind-mount patch) |
| `policy.megatron_cfg.optimizer.optimizer_cpu_offload` | False | False (tried True — made things worse) |
| `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` | n/a | **do not set** (vLLM `CuMemAllocator` asserts it out) |
I'm filing this as a single issue rather than splitting because the three
failures interact (free-memory math, vLLM sleep residue, and the IPC ping-pong
buffer all share the same 79 GB GPU), and the chronology is helpful for
anyone who hits them in this order. Three distinct upstream action items are
called out at the end.
### Environment
| Item | Value |
|---|---|
| NeMo-RL | v0.6.0 (container `nvcr.io/nvidia/nemo-rl:v0.6.0`) |
| Bundled Megatron-Bridge | `0.5.0+95e5f38f` |
| Bundled vLLM | `0.17.1` |
| PyTorch | bundled |
| Model | `Qwen/Qwen3.5-397B-A17B` (`text_config.mtp_use_dedicated_embeddings=False`, `num_experts=512`, `moe_intermediate=~1024`, `hidden=4096`) |
| Cluster | 32 × 8×H100 (79 GB visible per GPU) |
| Megatron parallelism | `TP=8`, `PP=8`, `EP=32`, `CP=1` |
| vLLM parallelism | `TP=16`, `EP` varied during debug (64 → 128 → 256) |
| Recipe | `examples/configs/recipes/llm/grpo-qwen3.5-397ba17b-32n8g-megatron.yaml` |
---
## Failure 1 — `ValueError: Object present on multiple PP ranks: [0, 7]`
### Symptom (initial run, lines 1066-1113 of `ray-driver.log`)
```
Traceback (most recent call last):
File "/opt/nemo-rl/examples/run_grpo.py", line 187, in
main()
File "/opt/nemo-rl/examples/run_grpo.py", line 113, in main
) = setup(config, tokenizer, dataset, val_dataset)
File "/opt/nemo-rl/nemo_rl/algorithms/grpo.py", line 747, in setup
state_dict_info = policy.prepare_refit_info()
File "/opt/nemo-rl/nemo_rl/models/policy/lm_policy.py", line 807, in prepare_refit_info
results = ray.get(futures)
ray.exceptions.RayTaskError(ValueError):
File "/opt/nemo-rl/nemo_rl/models/policy/workers/megatron_policy_worker.py", line 949, in prepare_refit_info
File "/opt/nemo-rl/nemo_rl/models/policy/workers/megatron_policy_worker.py", line 1003, in _calculate_refit_param_info
File "/opt/nemo-rl/nemo_rl/models/policy/workers/megatron_policy_worker.py", line 997, in calculate_size_in_bytes
File "/opt/nemo-rl/nemo_rl/models/megatron/pipeline_parallel.py", line 64, in broadcast_obj_from_pp_rank
raise ValueError(f"Object present on multiple PP ranks: {true_ranks}")
ValueError: Object present on multiple PP ranks: [0, 7]
```
The Ray `core_worker_process.cc:88` check failure that follows is a secondary
crash during driver teardown, not an independent bug.
### Root cause
1. `broadcast_obj_from_pp_rank` (`nemo_rl/models/megatron/pipeline_parallel.py:60-64`) uses `all_gather_object` to locate the single PP rank owning a given object, and **expects exactly one owner**. Two owners triggers the `[0, 7]` error.
2. `MegatronPolicyWorker._calculate_refit_param_info` calls `self.megatron_bridge.get_conversion_tasks([self.model])` to build `self.refit_conversion_tasks`. Each task carries a `param_weight` tensor that is non-None only on the owning PP rank — *for normal parameters*.
3. With **MTP shared embedding** (HF `mtp_use_dedicated_embeddings=False`), Megatron-LM places the MTP module on the last PP stage and **replicates** `embedding.word_embeddings.weight` onto that stage so MTP forward can reuse it.
4. `Megatron-Bridge.build_conversion_tasks` (`models/conversion/model_bridge.py:1407-1481`) only filters out `output_layer.weight` when `share_embeddings_and_output_weights=True` (`:1413`/`:1593`). It does **not** apply the analogous filter to the MTP-replicated `embedding.word_embeddings.weight` on the last PP stage.
5. Result: both PP=0 (canonical) and PP=last (MTP-replicated) hold a non-None `param_weight` for the same logical parameter; `broadcast_obj_from_pp_rank` reports `[0, 7]`.
Interesting upstream detail worth flagging: Megatron-Bridge **already defines**
a helper `_should_skip_mtp_duplicate_embedding_export`
(`model_bridge.py:1360`), and it is called from the HF-export path
(`:1135`). It is simply not wired into `build_conversion_tasks` (the path
NeMo-RL's refit uses).
### Workaround (carry-fix in NeMo-RL)
Inserted in `_calculate_refit_param_info`, immediately after the existing
`self.refit_conversion_tasks = [...]` assignment, before `param_info = []`:
```python
# Workaround for MTP shared-embedding double-ownership during refit.
#
# When the HF config sets `mtp_use_dedicated_embeddings=False` (e.g.
# Qwen3.5-MoE A17B), Megatron-LM places the MTP module on the last PP
# stage and replicates `embedding.word_embeddings.weight` there. Megatron-
# Bridge's `build_conversion_tasks` currently emits one task whose
# `param_weight` is non-None on BOTH PP=0 AND the last PP stage, which
# makes `broadcast_obj_from_pp_rank` raise `Object present on multiple PP
# ranks`. Drop the duplicate on the last PP stage; PP=0 stays the
# canonical owner.
from megatron.core.parallel_state import (
get_pipeline_model_parallel_world_size,
is_pipeline_last_stage,
)
if (
get_pipeline_model_parallel_world_size() > 1
and is_pipeline_last_stage(ignore_virtual=True)
):
for _t in self.refit_conversion_tasks:
if _t.param_name.endswith("embedding.word_embeddings.weight"):
print(
f"[mtp-shared-embed-skip] dropping duplicate on last PP "
f"stage: {_t.param_name}",
flush=True,
)
# WeightConversionTask is @dataclass(frozen=True);
# use object.__setattr__ to bypass the frozen check.
object.__setattr__(_t, "param_weight", None)
```
Delivered as a Pyxis single-file bind-mount on top of the read-only
container `.sqsh`; the launch script adds
`MOUNTS=$MOUNTS,:/opt/nemo-rl/.../megatron_policy_worker.py:ro`.
Verification line shows up on the last PP stage workers on each setup:
```
[mtp-shared-embed-skip] dropping duplicate on last PP stage: .embedding.word_embeddings.weight
```
(`` is empty or `language_model.` depending on bridge dispatched.)
### After Failure 1 cleared, `prepare_refit_info()` succeeds and GRPO loop reaches rollout.
---
## Failure 2 — IPC refit buffer too small for fused MoE experts
### Symptom (next run, ratio=0.3 default)
```
File "/opt/nemo-rl/nemo_rl/models/policy/utils.py", line 319,
in stream_weights_via_ipc_zmq_impl
assert aligned_size <= buffer_size_bytes, ...
AssertionError: Parameter model.language_model.layers.0.mlp.experts.gate_up_proj
too large for buffer: 8589934592 > 6263046144
```
### Math (corrected once during debug — see "Correction" below)
`refit_policy_generation` in `algorithms/grpo.py:1162-1169` sizes the IPC
buffer as `free_gpu_mem * NRL_REFIT_BUFFER_MEMORY_RATIO`, then `policy/utils.py:289`
halves it for ping-pong:
```python
buffer_size_bytes = buffer_size_bytes // 2
```
The assertion at `policy/utils.py:342` checks the **halved** value against
the largest fused parameter. So the real condition is:
```
(free_mem * ratio) / 2 >= largest_fused_param
```
Largest fused param on this model: `model.layers.*.mlp.experts.gate_up_proj`
at **8 GiB** (`8589934592 = 2³³`), fixed by:
```
num_experts(512) × 2 × moe_intermediate(~1024) × hidden(4096) × 2 bytes = 8 GiB
```
Crucially: **vLLM EP does not shrink this**. Megatron-Bridge
`_accumulate_grouped_export` (`model_bridge.py:815`) materializes the
global-experts HF tensor on the producer side via
`torch.stack([... for i in range(num_experts)])` where `num_experts = 512`
globally. vLLM slices its local experts from that global tensor on the
consumer side. Tensor size on the wire is fixed by model architecture, not
by vLLM's `expert_parallel_size`. (I confirmed this with a later
`VLLM_EP=128` run that produced the same 8 GiB error.)
### Tuning attempts and dead ends
Free memory at refit time, after Megatron's optimizer offload has run, is
**~38–44 GB** depending on vLLM EP (vLLM holds ~22–27 GB of weight residue
even in sleep mode level=1). I'll quote actual measured values below from
specific jobs.
#### Attempt 1: bump `NRL_REFIT_BUFFER_MEMORY_RATIO` 0.3 → 0.5
Buffer assertion passed at 0.5, but the **next** step OOMed during the EP
all-gather merge:
```
File ".../Megatron-Bridge/.../model_bridge.py", line 815, in _accumulate_grouped_export
merged = torch.stack([grouped_buffers[group_key][i] for i in range(num_experts)], dim=0)
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 8.00 GiB.
GPU 0: total 79.11 GB, free 6.76 GB.
Megatron worker process: 44.90 GB. vLLM process: 27.40 GB.
```
Buffer alloc and stack alloc both want 8 GiB. Larger IPC buffer → less free
at `stack` time.
#### Attempt 2: lower vLLM `gpu_memory_utilization` 0.6 → 0.5 (no effect)
Hoped to reduce vLLM's footprint and give the stack more room:
| | vLLM util 0.6 | vLLM util 0.5 |
|---|---|---|
| Megatron worker | 44.90 GB | 44.89 GB |
| vLLM residual | 27.40 GB | 27.34 GB |
| Free at OOM | 6.76 GB | 6.83 GB |
`gpu_memory_utilization` only sizes vLLM's **KV cache**, which is already
released by sleep mode level=1. The ~27 GB residual is per-GPU **weight
shard** (~397B / TP=16 / EP=64 ≈ 25-27 GB in bf16), invariant to that knob.
Reverted to 0.6 to preserve KV cache for rollout.
#### Attempt 3: enable `optimizer_cpu_offload=True` (made things worse)
Tried:
```yaml
policy.megatron_cfg.optimizer.optimizer_cpu_offload=True
policy.megatron_cfg.optimizer.optimizer_offload_fraction=1.0
```
Expected: free the ~3.3 GB of optimizer state Megatron leaves on GPU between
steps. Observed: Step 2 *after-offload* GPU residency went 12.72 GB (off) →
**15.86 GB (on)**. The offload feature appears to keep extra GPU buffers
(fp32 master / grad scaling staging) that the explicit `move_optimizer("cpu")`
path otherwise drops. Reverted.
#### Attempt 4: `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` (blocked by vLLM)
vLLM's `CuMemAllocator` (used for sleep-mode KV-cache release) asserts this
combination out at `cumem.py:132`:
```
AssertionError: Expandable segments are not compatible with memory pool.
Please track https://github.com/pytorch/pytorch/issues/147851 for the
latest updates.
```
There is no clean per-actor env override (each Ray actor inherits the
driver env), so this combination is unusable under vLLM's current
implementation. The launch script now carries a comment block explicitly
warning against re-adding the variable. **This is a real interoperability
gap** — flagging here in case maintainers want to surface the incompatibility
earlier (e.g., refuse to start if `PYTORCH_CUDA_ALLOC_CONF` contains
`expandable_segments:True` and vLLM is enabled).
#### Attempt 5: ratio=0.45 (fragmentation cliff)
After the ping-pong correction (see below), 0.45 was the smallest value
predicted to clear the per-side assertion (`(38.9 × 0.45) / 2 = 8.75 GB`).
Observed result was strictly *worse* than 0.5:
| Item | ratio=0.5 run | ratio=0.45 run | Δ |
|---|---|---|---|
| Megatron worker total | 44.89 GB | 48.69 GB | +3.80 GB |
| PyTorch allocated | 40.59 GB | 38.65 GB | -1.94 GB |
| **PyTorch fragmentation (reserved − allocated)** | **1.29 GB** | **7.04 GB** | **+5.75 GB** |
| Free at OOM | 6.83 GB | 2.91 GB | -3.92 GB |
The 8.75 GB ping-pong half does not pack cleanly alongside vLLM's ~27 GB
sleep-mode residue; the caching allocator reserves blocks larger than the
request and cannot reuse the leftovers for the subsequent 8 GB stack. Net:
~1.94 GB saved on IPC alloc was overwhelmed by +5.75 GB of fragmentation.
#### Correction during debug: the `// 2` ping-pong split was missed
My initial analysis used `free_mem * ratio` directly against the 8 GB
assertion. Reading the actual error messages backwards:
| ratio | per-side buffer in error | implied raw buffer | implied free_mem |
|---|---|---|---|
| 0.3 | 6.26 GB | 12.52 GB | ~41.7 GB |
| 0.4 | 7.78 GB | 15.56 GB | ~38.9 GB |
The assertion sees the **halved** value, not the raw. Corrected condition:
```
(free_mem * ratio) / 2 >= 8 GB
=> ratio >= 16 GB / free_mem ≈ 0.412 (with free ≈ 38.9 GB at EP=64)
```
So both 0.3 and 0.4 fail at EP=64 — consistent with the observed crashes.
#### Attempt 6: `VLLM_EP=128` (real benefit: vLLM residue, not fused tensor)
I had hypothesized doubling EP would halve the fused tensor. It did not
(see the architectural note above). What it **did** do was shrink each
vLLM worker's sleep-mode weight residue from ~27 GB to ~22 GB (4 local
experts instead of 8), pushing free_mem at refit time from ~38.9 GB to
~44 GB. That extra ~5 GB of headroom finally makes ratio=0.4 viable:
| Site | Formula | EP=128, ratio=0.4 | Margin |
|---|---|---|---|
| Assertion (per-side ≥ 8 GB) | `(44 × 0.4) / 2` | 8.8 GB | +0.8 GB |
| EP all-gather (free@stack ≥ 8 GB) | (extrapolated) | ~14 GB | +6 GB |
This was the first config that held through at least Step 1's refit. But
Step 2 was on the edge of the fragmentation cliff and inconsistent. The
final fix was a different lever — see Failure 3.
---
## Failure 3 — vLLM sleep residue + CUDA graph residue blocking Step 2+
After Failure 2 was tamed by `EP=128 + ratio=0.4`, runs still hit OOM at
**Step 2 refit** intermittently. The pattern was that **Step 1 refit ran
cleaner than Step 2** — counterintuitive, until I tracked what was actually
sitting on the GPU at each step's refit time:
- vLLM in sleep mode (level=1) holds the per-GPU weight shard (~22 GB at
EP=128, ~27 GB at EP=64). Level=1 releases only KV-cache-tagged
`CuMemAllocator` regions; weights and CUDA-graph capture buffers stay
put.
- vLLM with `enforce_eager=False` records CUDA graphs the first time it
runs each shape. By Step 2, those graphs have been recorded for both
prefill and decode shapes, adding ~10–15 GB of **non-cumem** buffers that
`CuMemAllocator.sleep()` cannot release at all.
So Step 2's refit sees less free memory than Step 1's, and once you're near
the edge it tips into OOM.
### Two coupled levers
1. **`policy.generation.vllm_cfg.enforce_eager=True`** — kills CUDA graph
capture entirely. Costs ~30–40% generation throughput, freed
~10–15 GB of non-cumem residue.
2. **vLLM `sleep(level=2)`** — discards everything at sleep, including the
weight shard, instead of backing it up to CPU. Since the next refit
re-streams weights from Megatron anyway, the CPU backup is wasted work
for colocated workflow. Delivered as a one-line bind-mount patch on
`vllm_worker.py`:
```python
# BEFORE
self.llm.sleep(level=1)
# AFTER (bind-mount patch)
self.llm.sleep(level=2)
```
The actual sleep_level=2 capability is in vLLM 0.17.1 (the bundled version)
— **no rebuild required**, just a parameter change.
Combined effect: vLLM-side sleep residue went from ~24 GB to ~7 GB. With
that recovered, IPC ratio can sit at **0.39** (per-side 8.6 GB at EP=256),
clear of the fragmentation cliff that bit ratio=0.45, with plenty of margin
at the `stack` site.
### Architectural note for upstream
Both of these would be safe and cheap to expose as config knobs in NeMo-RL
rather than requiring source patches:
- `policy.generation.vllm_cfg.sleep_level: int = 1` — pass through to
`self.llm.sleep(level=)` in `vllm_worker.py`. Default 1 preserves
current behavior.
- `enforce_eager` is already a `vllm_cfg` field, so no work needed.
Documenting *when* to set `sleep_level=2` (colocated + tight free_mem at
refit) would help users who hit this. The performance trade-off is bounded:
the CPU↔GPU weight backup at sleep takes a few seconds and is wasted in
colocated mode.
---
## Final verified config (smoke test pass, 2026-05-11)
Result: **9/10 GRPO steps fully passed** with KL Error in the 1.97–2.18
range, no OOM. Step 10 completed rollout + refit (no OOM at the previously-
failing 8 GiB stack) and was mid-training when the SLURM 4h time limit
fired (`STEP CANCELLED DUE TO TIME LIMIT`).
### Launch script (relevant excerpts from `nemo_rl_Qwen3_5_397B_A17B_Megatron.sh`)
```bash
# Cluster topology
NUM_ACTOR_NODES=32
TRAIN_TP=8
TRAIN_PP=8
TRAIN_EP=32
TRAIN_ETP=1
TRAIN_CP=1
VLLM_TP=16
VLLM_EP=256
# IPC refit buffer
export NRL_REFIT_BUFFER_MEMORY_RATIO=0.39
export PYTORCH_CUDA_ALLOC_CONF=garbage_collection_threshold:0.6
# Do NOT add expandable_segments:True — vLLM CuMemAllocator asserts it out
# (cumem.py:132; tracking pytorch/pytorch#147851).
# Hydra overrides on top of the recipe yaml
POLICY_GENERATION_ARGS=(
policy.generation.vllm_cfg.tensor_parallel_size=$VLLM_TP
policy.generation.vllm_cfg.expert_parallel_size=$VLLM_EP
policy.generation.vllm_cfg.enable_expert_parallel=True
policy.generation.vllm_cfg.max_model_len=$actor_ppo_max_token_len
policy.generation.vllm_cfg.enforce_eager=True
)
POLICY_MEGATRON_ARGS=(
policy.megatron_cfg.enabled=True
policy.megatron_cfg.tensor_model_parallel_size=$TRAIN_TP
policy.megatron_cfg.pipeline_model_parallel_size=$TRAIN_PP
policy.megatron_cfg.num_layers_in_first_pipeline_stage=6
policy.megatron_cfg.num_layers_in_last_pipeline_stage=6
policy.megatron_cfg.expert_model_parallel_size=$TRAIN_EP
policy.megatron_cfg.expert_tensor_parallel_size=$TRAIN_ETP
policy.megatron_cfg.context_parallel_size=$TRAIN_CP
# optimizer_cpu_offload=True made Step 2 GPU residency 12.72 -> 15.86 GB
# (extra fp32 master / grad scaling staging stayed on GPU).
# Do not re-enable without measurement.
)
# Bind-mount patches over the read-only container .sqsh
PATCH_FILES=(
"$RL_SCRIPT_DIR/megatron_policy_worker.py:/opt/nemo-rl/nemo_rl/models/policy/workers/megatron_policy_worker.py:ro"
"$RL_SCRIPT_DIR/vllm_worker.py:/opt/nemo-rl/nemo_rl/models/generation/vllm/vllm_worker.py:ro"
)
for p in "${PATCH_FILES[@]}"; do MOUNTS="${MOUNTS},${p}"; done
# SLURM
#SBATCH --time=04:00:00
```
### Bind-mounted files (md5 of the working set)
| File | md5 | What it contains |
|---|---|---|
| `nemo_rl_Qwen3_5_397B_A17B_Megatron.sh` | `73a8a7b59d7ee4920c30723a3855b887` | Launch script as above |
| `megatron_policy_worker.py` | `5eeeee0177cf8689c602eb4ac1817378` | MTP shared-embedding refit fix (Failure 1 workaround) |
| `vllm_worker.py` | `67ef27eedab6d09372c2947750b01cbd` | `self.llm.sleep(level=2)` (Failure 3 lever) |
(Backed up at `RL_SCRIPT/*.bak.smoke_pass_20260510` on the reporter's
filesystem.)
### Per-step measurements at the smoke pass
(Steps 1-9 — Step 10 was cut by SLURM time limit, not by any of the issues
above. Numbers approximate; aggregated from `ray-driver.log`.)
- IPC refit assertion (per-side): pass with ~0.6 GB margin
- `stack` site at `_accumulate_grouped_export`: pass with several GB margin
- vLLM sleep residue: ~7 GB (was ~24 GB pre-fix)
- KL Error: 1.97 – 2.18 across steps
---
## Bonus observation: dynamic-sampling fill rate is hard to read mid-run
(Not strictly part of the same bug, but I hit this in an immediate follow-up
run that re-enabled `grpo.use_dynamic_sampling=True` on top of the
smoke-pass config, and it relates to the same workflow. Filing here for
visibility; happy to spin off a separate issue if maintainers prefer.)
With the smoke-pass config + `grpo.use_dynamic_sampling=True` + recipe defaults
(`batch_multiplier=2`, `dynamic_sampling_max_gen_batches=10`,
`train_global_batch_size=256`), Step 1 took 3 gen batches to fill the
buffer (cumulative 0 → 184 → 344 non-zero-std prompts), wall time **4561 s
(76 min)**, of which 82.5% was generation. Measured non-zero-std hit rate:
**~18%** per generated sample. On this base policy + DAPO math, most
prompts produce identical-reward rollouts (mostly all-wrong); the
data-policy mismatch is real, not a cold-start artifact.
Three minor improvements to dynamic sampling observability would have saved
me time:
1. `grpo.py:891-894` prints only the **cumulative** non-zero-std buffer size.
Adding the per-batch delta + hit rate would let users predict whether
`max_gen_batches` will be hit within 1-2 gen batches instead of waiting
for the full sequence.
2. `grpo.py:910` raises `ValueError` on `max_gen_batches` exhaustion with
a generic message ("Consider evaluating the complexity of your data or
adjusting the num_prompts_per_step or num_generations_per_prompt"). It
doesn't quote the measured fill rate or point at the most-impactful
levers (`train_global_batch_size`, `dynamic_sampling_max_gen_batches`,
`batch_multiplier`).
3. The metrics aggregator at `grpo.py:1932-1935` silently drops scalar
`int` / `float` metrics including `dynamic_sampling_num_gen_batches`
and `dynamic_sampling_num_discarded_valid_samples` (log line
`Skipping aggregation for X ()`). These are the most
important dynamic-sampling telemetry. Adding an `isinstance(v, (int, float))`
passthrough branch fixes it.
---
## Suggested upstream actions
In rough order of impact, listing as observations rather than committing to
specific patches (happy to discuss the right surface for any of these):
### High impact
1. **Megatron-Bridge: wire `_should_skip_mtp_duplicate_embedding_export`
into `build_conversion_tasks`.** The helper already exists
(`model_bridge.py:1360`) and is used in the HF-export path (`:1135`).
The refit path that NeMo-RL uses (`build_conversion_tasks`, around
`:1407-1481`) does not call it. Without this, any MTP-shared-embedding
model with Megatron PP > 1 fails refit at setup. Suggested issue title
for the Megatron-Bridge side:
> `build_conversion_tasks` does not skip MTP-replicated word embedding on last PP stage when `mtp_use_dedicated_embeddings=False`, causing "Object present on multiple PP ranks" during refit on Qwen3.5-MoE with PP > 1
2. **NeMo-RL: defensive carry-fix for the above.** Until Megatron-Bridge
lands the upstream fix and NeMo-RL bumps the pin, the 6-line guard in
`_calculate_refit_param_info` quoted above would let users run today.
Happy to send a small PR if there's appetite.
3. **NeMo-RL: expose vLLM sleep level as config.** Currently hardcoded to
`level=1` in `vllm_worker.py::sleep`. Adding
`policy.generation.vllm_cfg.sleep_level: int = 1` and passing it
through would let users opt into `level=2` for colocated workflows
without bind-mounting. Documentation pointer ("set level=2 if you see
tight free memory at refit and weights will be re-streamed anyway")
would also help.
### Medium impact
4. **NeMo-RL: surface the `expandable_segments:True` vs vLLM
incompatibility early.** Currently fails ~3-5 minutes into setup at
`cumem.py:132`. A startup-time check (refuse to launch if
`PYTORCH_CUDA_ALLOC_CONF` contains `expandable_segments:True` and vLLM
is colocated, or downgrade with a warning) would save users a debug
loop.
5. **NeMo-RL: dynamic-sampling observability** (the three bullets above).
Pure additive observability, no behavior change.
### Lower impact / docs
6. **Document the IPC refit buffer math.** The interaction between
`NRL_REFIT_BUFFER_MEMORY_RATIO`, the `// 2` ping-pong split, the EP
all-gather staging, and the fixed 8 GiB fused-MoE tensor isn't
obvious from reading the code. A short note in the GRPO config docs
pointing at `policy/utils.py:289` and `algorithms/grpo.py:1162-1169`
would help users tune intentionally instead of binary-searching
ratios.
7. **Document the `optimizer_cpu_offload=True` cost on Megatron MoE.**
Counterintuitive that enabling offload *increases* post-step GPU
residency by ~3 GB (fp32 master / grad scaling staging stays on GPU).
---
## Reproducibility / artifacts
All numbers above come from a sequence of runs the reporter executed over
2026-05-08 to 2026-05-12, summarized in chronological order below:
| Run | What it shows |
|---|---|
| Run A | Original failure (MTP shared-embedding ValueError) |
| Run B | Failure 2 — ratio=0.3 IPC assertion |
| Run C | Failure 2 — ratio=0.5 EP all-gather OOM |
| Run D | Failure 2 — ratio=0.4 IPC assertion (correction discovered here) |
| Run E | ratio=0.45 fragmentation cliff |
| Run F | VLLM_EP=128 fused tensor still 8 GiB |
| Run G | optimizer_cpu_offload=True regression |
| Run H | enforce_eager + sleep level=2 stabilization |
| **Run I** | **Smoke test pass (9/10 steps)** |
| Run J | Dynamic-sampling re-validation (bonus observation) |
Happy to share specific log excerpts on request.
Contributor guide
Assessment
This issue has not been assessed yet.