NVIDIA-NeMo / NVIDIA-NeMo/RL

Qwen3.5 Automodel path is slower than MCore because flash-linear-attention (FLA) is not installed by default

Open
#2,324 5 comments 0 reactions 1 assignee Claimed by @ashors1 View on GitHub
bug Documentation Performance
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

# Summary

On the Automodel (FSDP2) path, Qwen3.5 models run measurably slower than the MCore path because `flash-linear-attention` (FLA) is not installed in the Automodel worker venv by default. Several kernels in `nemo-automodel` dispatch to FLA when it is importable and fall back to slower implementations otherwise.

The root-cause fix belongs in `nemo-automodel` and is tracked separately — see https://github.com/NVIDIA-NeMo/Automodel/pull/1894, which moves FLA out of the dev-only dependency group so downstream installers (including NeMo-RL's `automodel` extra) pick it up transitively via `nemo-automodel[moe]`.

This issue tracks the NeMo-RL side:

1. Document the symptom so users picking the automodel path for Qwen3.5 (dense or MoE) are aware.
2. Bump `3rdparty/Automodel-workspace/Automodel` ; after that, no manual workaround should be needed.

# Symptoms

Same workload (Qwen3.5-35B-A3B-Base, SFT, BF16, 4K seq, 2n8g), three setups:

- Yellow: MCore path
- Red: Automodel + FLA installed
- Blue: Automodel without FLA (default today)

Image

# Affected recipes (NeMo-RL side)

Automodel-path recipes for Qwen3.5 dense and MoE, e.g.:

- `examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-2n8g-automodel-ep16.yaml`
- `examples/configs/recipes/llm/grpo-qwen3.5-35ba3b-dapo-4n8g-automodel.yaml`
- `examples/configs/recipes/vlm/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-automodel-ep16.yaml`

(`-megatron` recipes use MCore kernels directly and are not affected.)

Also note: CP on Qwen3.5 MoE via the automodel path explicitly requires FLA today — see the `import fla` guard in `nemo_rl/models/automodel/setup.py`.

# Resolution criteria

- [x] Upstream: https://github.com/NVIDIA-NeMo/Automodel/pull/1894 merges (moves FLA to optional extras so `nemo-automodel[moe]` pulls it in by default).
- [ ] NeMo-RL: bump `3rdparty/Automodel-workspace/Automodel` to a commit that includes #1894; after that, the Automodel worker venv will install FLA automatically and this issue can be closed.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.