[NeMo-RL-DSV4] Review peak memory from FP32 cast before MoE grouped-mm scatter_add

During DeepSeek V4 GRPO bring-up in NeMo-RL, we observed an OOM in the Automodel policy path when using backend.experts: torch_mm. We have not reproduced this issue in a standalone Automodel training run, so this report is mainly to share the integration-time observation and ask for Automodel maintainers' review.

Issue

The suspected peak-memory source is in the grouped MoE forward path. After the down projection, Automodel accumulates expert outputs into an FP32 buffer via:

scatter_ids = sorted_token_ids.unsqueeze(1).expand_as(output2)
y.scatter_add_(0, scatter_ids, output2.float())

For large MoE / GRPO shapes, output2.float() appears to materialize a full FP32 temporary tensor before the scatter, which can create a large transient memory spike.

Local workaround in fork

In my fork, I added a local mitigation in this commit:

https://github.com/zpqiu/Automodel/commit/e2564e22c107369db814d57de71dfa42c5b5d047

The change introduces an optional env var:

NEMO_AUTOMODEL_MOE_SCATTER_CHUNK_ROWS

When unset or 0, behavior stays unchanged. When set, the FP32 cast and scatter are done in row chunks:

for start in range(0, output.size(0), chunk_rows):
    end = min(start + chunk_rows, output.size(0))
    y.scatter_add_(0, scatter_ids[start:end], output[start:end].float())

This is intended to preserve FP32 accumulation semantics while reducing the peak temporary buffer size from the full output2.float() tensor to one chunk at a time. In the NeMo-RL DS V4 GRPO experiment path, I used NEMO_AUTOMODEL_MOE_SCATTER_CHUNK_ROWS=4096 as a temporary workaround.

Request

Could you review whether this memory pattern is expected in Automodel's grouped MoE path, and whether a chunked scatter approach, or another upstream-friendly implementation, makes sense?

This workaround has not been rigorously validated for standalone Automodel correctness or performance.

Start by reviewing the grouped MoE forward path containing the FP32 cast and scatter_add_, then compare it with the local workaround in commit e2564e22c107369db814d57de71dfa42c5b5d047. Validate the chunked approach in standalone Automodel runs for correctness, performance, and peak memory, including the NEMO_AUTOMODEL_MOE_SCATTER_CHUNK_ROWS=4096 setting.