[Bug] VIME: Gemma4-26B-MOE, Garbled Rollout Outputs on 4×8 A800 40GB GPUs with EP4/PP2/TP2 ,Notice that HF1 !=HF1-Megatron-HF2

[chenchao9999]

Bug Description

Environment

• Framework: VIME
• GPU: NVIDIA A800 40GB
• Cluster: 4 nodes × 8 GPUs
• Parallelism: EP4 / PP2 / TP2

Problem

The rollout outputs are garbled.

Preliminary investigation suggests that the issue may be related to weight transfer/synchronization between the actor and rollout workers.

The torch_dist checkpoint was generated using the official VIME conversion script.

Training / Inference Deployment

I am using the disaggregated training and inference mode, with actor and rollout workers separated.

The relevant configuration is:

--actor-num-nodes "${NNODES}" \
--actor-num-gpus-per-node "${ACTOR_GPUS_PER_NODE}" \
--rollout-num-gpus "${ROLLOUT_NUM_GPUS}" \

### Steps to Reproduce

# VIME: Garbled Rollout Outputs on 4×8 A800 40GB GPUs with EP4/PP2/TP2

## Environment

* Framework: VIME
* GPU: NVIDIA A800 40GB
* Cluster: 4 nodes × 8 GPUs
* Model parallelism: EP4 / PP2 / TP2
* Training/inference mode: Disaggregated actor and rollout deployment

## Problem

The rollout outputs are garbled.

Preliminary investigation suggests that the issue may be related to the **weight transfer/synchronization process** between the actor and rollout workers.

The initial `torch_dist` checkpoint was generated using the **official VIME conversion script**.

## Training / Inference Deployment

I am using a disaggregated training and inference setup, with actor and rollout workers separated.

```bash
--actor-num-nodes "${NNODES}" \
--actor-num-gpus-per-node "${ACTOR_GPUS_PER_NODE}" \
--rollout-num-gpus "${ROLLOUT_NUM_GPUS}" \

The deployment configuration is:

• Actor: 4
• Rollout: 4
• TP = 2
• PP = 2
• EP = 4

Key Configuration

Checkpoint

CKPT_ARGS=(
    --hf-checkpoint "$HF_CKPT"
    --ref-load "$_REF_PATH"
    --load "$_LOAD_PATH"
    --save "$SAVE_DIR"
    --save-interval "$SAVE_INTERVAL"
)

Rollout

ROLLOUT_ARGS=(
    --prompt-data "$PROMPT_DATA"
    --input-key prompt
    --label-key label
    --apply-chat-template
    --apply-chat-template-kwargs '{"enable_thinking": false}'
    --rollout-shuffle
    --group-rm
    --custom-rm-path rewards.translation_rm.reward_func
    --rollout-function-path rollout.gen_then_batch_rm.generate_rollout
    --num-rollout "$NUM_ROLLOUT"
    --rollout-batch-size "$ROLLOUT_BS"
    --over-sampling-batch-size "$OVER_SAMPLING_BS"
    --n-samples-per-prompt "$N_SAMPLES"
    --rollout-max-response-len "$MAX_RESP"
    --rollout-temperature "$TEMPERATURE"
    --num-steps-per-rollout "$NUM_STEPS_PER_ROLLOUT"
    --global-batch-size "$GLOBAL_BS"
    --balance-data
)

Model Parallelism / Performance

PERF_ARGS=(
    --tensor-model-parallel-size "$TP"
    --sequence-parallel
    --pipeline-model-parallel-size "$PP"
    --context-parallel-size "$CP"
    --expert-model-parallel-size "$EP"
    --expert-tensor-parallel-size "$ETP"
    --recompute-granularity full
    --recompute-method uniform
    --recompute-num-layers 1
    --use-dynamic-batch-size
    --calculate-per-token-loss
    --max-tokens-per-gpu "$MAX_TOKENS_PER_GPU"
)

GRPO

GRPO_ARGS=(
    --advantage-estimator grpo
    --use-kl-loss
    --kl-loss-coef 0.00
    --kl-loss-type low_var_kl
    --entropy-coef 0.00
    --eps-clip 0.2
    --eps-clip-high 0.28
)

Optimizer

OPTIMIZER_ARGS=(
    --optimizer adam
    --lr "$LR"
    --lr-decay-style constant
    --weight-decay 0.1
    --adam-beta1 0.9
    --adam-beta2 0.98
    --optimizer-cpu-offload
    --overlap-cpu-optimizer-d2h-h2d
    --use-precision-aware-optimizer
)

vLLM Rollout

VLLM_ARGS=(
    --rollout-num-gpus-per-engine "$VLLM_TP"
    --vllm-gpu-memory-utilization "${VLLM_GPU_MEMORY_UTILIZATION}"
    --vllm-max-model-len "$((MAX_RESP + 2048))"
    --vllm-max-cudagraph-capture-size "${VLLM_MAX_CUDAGRAPH_CAPTURE_SIZE:-1}"
    --vllm-max-num-seqs "${VLLM_MAX_NUM_SEQS:-32}"
    --vllm-server-concurrency "${VLLM_SERVER_CONCURRENCY:-128}"
    --vllm-language-model-only
)

Weight Update / Synchronization

MISC_ARGS=(
    --attention-dropout 0.0
    --hidden-dropout 0.0
    --accumulate-allreduce-grads-in-fp32
    --attention-softmax-in-fp32
    --attention-backend flash
    --loss-mask-type gemma4
    --megatron-to-hf-mode "${MEGATRON_TO_HF_MODE}"
    --check-weight-update-equal
    --update-weight-mode full
    --update-weight-transport disk
    --update-weight-disk-dir /xx/vime_output/gemma4_weight_disk_debug
    --update-weight-disk-keep-files
)

In particular, I am using full weight updates with disk-based transport:

--update-weight-mode full
--update-weight-transport disk
--check-weight-update-equal

Current Investigation

The initial torch_dist checkpoint was generated using the official VIME conversion script.

At the moment, the issue appears to be related to the weight migration/update path, potentially involving one of the following stages:

1. torch_dist checkpoint loading on the actor side.
2. Megatron-to-HF weight conversion.
3. Actor-to-rollout weight transfer.
4. Weight synchronization on the vLLM rollout workers.

The rollout model can start successfully, but the generated outputs are garbled, which suggests that the rollout workers may not be receiving or reconstructing the correct model weights.

I am currently using:

--megatron-to-hf-mode "${MEGATRON_TO_HF_MODE}"
--check-weight-update-equal
--update-weight-mode full
--update-weight-transport disk

Any suggestions on how to further verify the correctness of the weight conversion and actor-to-rollout synchronization would be appreciated.

Expected Behavior

Additional Investigation

I expected the original HF checkpoint (HF1) and the checkpoint obtained through the HF1 -> Megatron -> HF2 conversion process to produce consistent inference results when deployed with vLLM.

Actual Behavior

However, the actual results are:

• HF1: normal outputs
• HF2: abnormal / garbled outputs

Both checkpoints were tested independently using vLLM.

Therefore, the issue appears to occur during the following conversion path:

HF1 -> Megatron (torch_dist) -> HF2

### Environment

- vime version: latest branch, commit `6cefd8463f7fa6b80b6692175975d91975f6bd3c`
- Python version: from the latest official VIME Docker image
- PyTorch version: from the latest official VIME Docker image
- CUDA/ROCm version: from the latest official VIME Docker image
- GPU type and count: 4 nodes × 8 NVIDIA A800 40GB GPUs (32 GPUs total)
- OS: Linux
- vLLM version: from the latest official VIME Docker image
- vllm-router version: from the latest official VIME Docker image
- Megatron-LM version (if relevant): from the latest official VIME Docker image

### Logs

```shell

Additional Context

Runtime Behavior

The training logs look normal and no obvious errors are reported.

However, the rollout outputs are garbled/abnormal.

In other words:

• Training process starts and runs normally.
• No obvious exceptions are observed in the logs.
• Actor/rollout processes appear to initialize successfully.
• The rollout generation itself is abnormal and produces garbled text.

Combined with the previous test:

• HF1 deployed with vLLM: normal output
• HF1 -> Megatron -> HF2, then deploy HF2 with vLLM: garbled output

This suggests that the issue is likely related to the model weight conversion path rather than a runtime crash or an obvious rollout initialization failure.

Pre-submission Checklist

• I have read the CONTRIBUTING.md and understand the collaboration scope.
• I have read the documentation and my issue is not addressed there.
• I have searched for existing issues and this is not a duplicate.
• I have provided a minimal, reproducible example.

[chenchao9999]

MEGATRON_TO_HF_MODE == raw

[chenchao9999]

我的探针排查链路

当前问题总结

1. 问题现象

当前权重转换链路为：

Megatron checkpoint
→ Megatron model / local weights
→ EP / PP / TP 权重重建
→ Megatron → HF Gemma4 converter
→ HF safetensors
→ vLLM load
→ rollout

目前现象：

• 原始 HF 权重使用 vLLM 推理正常。
• HF → Megatron → HF 后保存出的 HF 权重使用 vLLM 推理出现乱码。
• 对 GOOD HF 和 BAD HF 的全部 30 层 MoE expert 权重进行比较后发现：
  • 30 层全部 mismatch。
  • expert 权重 cosine 大约只有 0.24 ~ 0.26。
• 因此问题基本可以确定在 Megatron → HF 的 MoE expert 权重恢复/同步链路，而不是 tokenizer、chat template、sampling 或 vLLM 本身。

---

2. 当前最怀疑：EP expert reconstruction / broadcast

当前训练拓扑：

• TP = 2
• PP = 2
• EP = 4
• CP = 1
• MP = 16
• Actor GPU = 16
• DP = 1

EP=4 时，可以抽象理解为：

• EP rank 0 → experts 0~31
• EP rank 1 → experts 32~63
• EP rank 2 → experts 64~95
• EP rank 3 → experts 96~127

Megatron → HF 时必须重新恢复成完整的：

expert 0
expert 1
...
expert 127

之前的探针发现了非常明显的异常：

• expert 0 和 expert 32 出现完全相同的 fingerprint。
• expert 31 和 expert 127 出现完全相同的 fingerprint。
• 类似的重复一直传递到了 BUFFER 和 STACK。

这种现象不像 BF16 精度误差或者 safetensors 保存误差，更像是：

1. EP broadcast 的通信顺序不一致；
2. expert owner / src_rank 映射错误；
3. global expert ID 和 local expert ID 映射错误；
4. 某个 EP rank 的 expert tensor 被 broadcast 到了错误的 expert 位置。

因此当前第一嫌疑是：

_get_megatron_full_params() 中的 EP expert reconstruction / broadcast。

---

3. 第二种可能：megatron_local_weights 本身已经错误

因此新增 LocalOwnerProbe，在任何 PP / EP / TP 通信之前直接检查：

megatron_local_weights[info.name]

当前完整探针链路：

Megatron model
↓
megatron_local_weights
↓
LocalOwner
↓
SOURCE
↓
PP broadcast
↓
EP reconstruction
↓
EPFULL
↓
TP all_gather
↓
FULL
↓
Gemma4 converter
↓
BUFFER
↓
STACK
↓
HF

如果 LocalOwner 中：

• expert 0 = A
• expert 32 = B
• expert 64 = C
• expert 96 = D

并且：

A != B != C != D

说明 Megatron 本地 expert 权重是正常的。

如果随后 EPFULL 变成：

• expert 0 = A
• expert 32 = A
• expert 64 = A
• expert 96 = A

则可以直接确定：

EP reconstruction / broadcast 把 expert 权重覆盖或串位了。

这是目前认为概率最高的情况。

---

4. 如果 LocalOwner 就已经错误

如果在任何通信之前就出现：

• expert 0 = A
• expert 32 = A
• expert 64 = A
• expert 96 = A

那么 EP broadcast 就不是根因。

需要继续向前检查：

Megatron checkpoint
↓
checkpoint load
↓
named_params_and_buffers
↓
ParamInfo
↓
src_rank
↓
megatron_local_weights

这时最可疑的是：

• global expert ID 与 local expert ID 的对应关系；
• ParamInfo.name；
• ParamInfo.src_rank；
• EP rank 与 expert owner 的映射。

---

5. 第三种可能：TP all_gather

如果探针结果为：

• LocalOwner：GOOD
• SOURCE：GOOD
• EPFULL：GOOD
• FULL：BAD

则问题发生在：

all_gather_params_async()

也就是：

EP reconstruction
↓
EPFULL：正常
↓
TP all_gather
↓
FULL：异常

这时需要重点检查：

• expert tensor parallel group；
• regular tensor parallel group；
• expert_tensor_parallel_world_size；
• tensor_model_parallel_world_size；
• expert 参数在 TP gather 中的拼接维度和顺序。

目前该可能性低于 EP reconstruction。

---

6. Gemma4 converter 当前不是第一嫌疑

Gemma4 converter 中 expert 的处理逻辑主要是：

bucket[expert_idx] = param

然后：

expert 0 ... expert 127

按照 expert index 排序后：

torch.stack(...)

最终生成：

• experts.gate_up_proj
• experts.down_proj

之前探针已经发现：

BUFFER 阶段 expert tensor 就已经发生重复。

因此如果：

BUFFER expert0 == BUFFER expert32

那么后面的 STACK：

STACK expert0 == STACK expert32

是必然结果。

所以当前更像是：

Gemma4 converter 正确地 stack 了上游已经错误的 expert tensor。

因此暂时没有必要继续修改 gemma4.py 的 expert stack 逻辑。

---

7. tied embedding 是另一个独立问题

Gemma4 配置：

tie_word_embeddings = True

converter 中：

embedding.word_embeddings.weight
→ model.language_model.embed_tokens.weight

同时：

output_layer.weight
→ model.language_model.embed_tokens.weight

因此之前出现：

Duplicate HF tensor while finalizing: model.language_model.embed_tokens.weight

这个问题是真实存在的。

但探针已经确认：

• embedding hash
• output_layer hash

两者完全一致。

所以 tied embedding duplicate 是一个需要修复的独立问题，但它无法解释：

全部 30 层 expert 权重 cosine 只有约 0.25。

因此它不是当前 rollout 乱码的主要嫌疑。

---

8. 当前问题概率排序

位置	可能性	判断依据
EP expert reconstruction / broadcast	最高	出现明显的跨 EP shard expert 复制/串位
ParamInfo / src_rank / global expert ID 映射	高	如果 LocalOwner 已错误，则重点检查这里
TP all_gather_params_async	中	EPFULL 正常但 FULL 异常时成立
Gemma4 BUFFER / STACK	低	当前看起来是在传播上游错误
tied embedding	独立问题	会造成 duplicate，但无法解释全部 expert mismatch
HF shard saver	很低	主要负责 tensor 写盘和 index
vLLM	很低	HF 权重离线比较已经证明 expert 权重不一致
tokenizer / chat template	基本排除	无法造成 expert cosine ≈ 0.25

---

9. 当前定位策略

现在不继续猜具体是哪段代码，而是寻找：

expert tensor 第一次发生错误的位置。

判断路径：

LocalOwner
├── BAD
│ └── 检查 ParamInfo / src_rank / Megatron local expert mapping
│
└── GOOD
↓
SOURCE
↓
EPFULL
├── BAD
│ └── EP reconstruction / broadcast BUG（目前最怀疑）
│
└── GOOD
↓
FULL
├── BAD
│ └── TP all_gather BUG
│
└── GOOD
↓
BUFFER / STACK
└── BAD → Gemma4 converter BUG

因此下一次运行最重要的是比较：

• LocalOwner
• SOURCE
• EPFULL
• FULL
• BUFFER
• STACK

尤其关注 EP shard 边界 expert：

• 0 / 31
• 32 / 63
• 64 / 95
• 96 / 127

只要确定这些 expert 第一次在哪个阶段从不同 fingerprint 变成相同 fingerprint，基本就可以把问题定位到具体函数。

[aoshen02]

I traced the reported corruption to the old raw Megatron-to-HF EP reconstruction path.

The issue revision (6cefd846) predates #340. In that revision, _get_megatron_full_params() passed a global ParamInfo.src_rank into an EP subgroup and fell back to each caller's own rank when the owner was outside that subgroup. With PP2/EP4, subgroup members could therefore disagree on the broadcast root. This matches the probe result here: experts copied across EP boundaries (0/32 and 31/127) before BUFFER/STACK. #340 replaced this with an explicit global-owner -> EP-group-local-root map and has regression tests for that topology.

PR #386 subsequently removed the old custom Gemma4 implementation, so I restored Gemma4 using current Megatron-Bridge rather than reviving that stack: aoshen02#2 (base is the #386 head branch).

Validation used Gemma4-26B-A4B-it on 2x8 H200 with the same TP2/PP2/EP4 topology (ETP2):

• HF -> Megatron -> HF completed on all 16 ranks
• 657/657 text tensors were bitwise exact
• no missing/extra/shape/dtype differences
• experts 0/31/32/63/64/95/96/127 all matched exactly
• separate vLLM runs on HF1 and reconstructed HF2 produced identical token IDs for three greedy prompts

I did not rerun the full 4-node actor/rollout deployment, but the independently isolated 16-rank conversion boundary that produced the bad HF2 checkpoint is covered directly.