Dev (Qwen 3.5 recipe)

README.md

@@ -43,6 +43,7 @@ AstraFlow **natively** supports the following for LLM RL training **without any
 <!-- <p align="center"><i>AstraFlow training a multi-policy workflow on an elastic, heterogeneous, cross-region rollout pool — all at once, with no feature-specific code.</i></p> -->
 
 ## News
+- **[2026/07]** AstraFlow **v0.1.2** images released — transformers 5.8 / SGLang 0.5.13 stack, new **Qwen3.5-4B (Gated-DeltaNet)** math recipes, Megatron-image build fix. Both images validated end-to-end on 8×H100 (400-step math-RL runs incl. eval). See `docker/README.md`.
 - **[2026/06]** New recipe: **dynamic recursive agent** on TextCraft — a multi-turn agent that recursively spawns sub-agents sharing inventory under a team reward. See the [recipe docs](https://Infini-AI-Lab.github.io/astraflow/docs/recipes/textcraft-recursive.html).
 - **[2026/06]** AstraFlow **v0.1.1** released — CUDA 13 image, SGLang 0.5.12, Megatron weight-sync training backend, and transformers 5 support. See the [project website](https://Infini-AI-Lab.github.io/astraflow/).
 - **[2026/05]** AstraFlow **v0.1.0** released — first public release of the full system. See the [project website](https://Infini-AI-Lab.github.io/astraflow/).

astraflow/train_worker/api/cli_args.py

@@ -473,10 +473,22 @@ class TrainEngineConfig:
     trial_name: str = ""
     path: str = field(default="", metadata={"help": "Path to HuggingFace checkpoint"})
     attn_impl: str = field(
-        default="flash_attention_2",
+        default="kernels-community/flash-attn2",
         metadata={
-            "help": "Attention implementation for huggingface transformers model.",
-            "choices": ["flash_attention_2"],
+            "help": (
+                "Attention implementation for huggingface transformers model. "
+                "Default pulls a prebuilt FlashAttention-2 kernel from the HF kernels "
+                "hub (ABI-matched to torch, incl. varlen for packed sequences). The "
+                "literal 'flash_attention_2' loads the local flash-attn wheel, which is "
+                "ABI-broken on torch>=2.11+cu13; 'sdpa' works but relies on position_ids "
+                "resets for packed block-diagonal masking."
+            ),
+            "choices": [
+                "kernels-community/flash-attn2",
+                "flash_attention_2",
+                "sdpa",
+                "eager",
+            ],
         },
     )
     init_from_scratch: bool = field(

astraflow/train_worker/engine/fsdp_engine.py

@@ -94,8 +94,6 @@
 from astraflow.train_worker.utils.model import (
     disable_dropout_in_model,
     is_gemma3_model,
-    is_qwen3_moe_model,
-    is_qwen3_vl_model,
     is_qwen_vl_model,
     is_valid_vision_model,
 )
@@ -1206,16 +1204,15 @@ def _prepare_mb_list(self, input_: dict[str, Any]) -> MicroBatchList:
             ]
             mb["use_cache"] = False
             padded_mb["use_cache"] = False
-            if is_qwen3_moe_model(self.model_config.model_type) or is_qwen3_vl_model(
-                self.model_config.model_type
-            ):
-                mb["attention_mask"] = None
-                padded_mb["attention_mask"] = None
-            else:
-                mb["attention_mask"] = dict(full_attention=None, sliding_attention=None)
-                padded_mb["attention_mask"] = dict(
-                    full_attention=None, sliding_attention=None
-                )
+            # Always pass attention_mask=None for the packed/varlen forward: per-sequence
+            # causal masking is driven by cu_seqlens + position_ids, and the model builds
+            # the right mask from None. The old dict(full_attention=None,
+            # sliding_attention=None) form is a transformers-4.x relic: on transformers>=5
+            # a dense model (qwen3 / qwen2) treats that dict as a *precomputed* mask, skips
+            # creation, and crashes. Passing None lets the model build its mask from
+            # cu_seqlens + position_ids instead.
+            mb["attention_mask"] = None
+            padded_mb["attention_mask"] = None
             if "multi_modal_input" in mb:
                 image_grid_thw_list = [
                     item["image_grid_thw"]

astraflow/train_worker/trainer/ppo_trainer.py

@@ -543,6 +543,33 @@ def train(
                 and getattr(config.evaluator, "freq_steps", None) is not None
             ):
                 if self._is_rank0:
+                    # Gate the pre-train eval on RaaS-pool readiness. The trainer
+                    # reaches loop entry tens of seconds before the SGLang
+                    # inference servers finish starting and register with the
+                    # dataflow pool; firing the eval immediately hits
+                    # "no healthy RaaS instance available for eval" and the v=0
+                    # baseline is silently dropped. Poll /status until at least
+                    # one RaaS instance is registered, then let init settle.
+                    import time as _time
+
+                    _deadline = _time.monotonic() + 300.0
+                    _raas_ready = False
+                    while _time.monotonic() < _deadline:
+                        try:
+                            _pool = self.astraflow.get_status().get("raas_pool", {})
+                            if int(_pool.get("size", 0)) >= 1:
+                                _raas_ready = True
+                                break
+                        except Exception:
+                            pass
+                        _time.sleep(2.0)
+                    if _raas_ready:
+                        _time.sleep(5.0)  # let the just-registered engine finish init
+                    else:
+                        logger.warning(
+                            "[pre-train] RaaS pool not ready after 300s; "
+                            "initial eval may be skipped",
+                        )
                     print(
                         "[Trainer] [pre-train] Running eval at version=0 ...",
                         flush=True,

astraflow/train_worker/utils/model.py

@@ -5,6 +5,9 @@
     "qwen2_vl",
     "qwen2_5_vl",
     "qwen3_vl",
+    # Qwen3.5 dense math checkpoints ship as Qwen3_5ForConditionalGeneration, so they
+    # load via the ImageTextToText path even though these recipes train text-only.
+    "qwen3_5",
     "gemma3",
 ]
 # Registry of vision models verified to work with this framework.

docker/Dockerfile.sglang

@@ -29,8 +29,14 @@ RUN uv pip install -e .
 # --- sglang layers (on top of basic) ---
 RUN uv pip install -e ".[sglang]"
 
-# flash-attn is excluded from uv resolution (see pyproject.toml [tool.uv]),
-# so install it explicitly with --no-build-isolation.
-RUN uv pip install "flash-attn==2.8.3" --no-build-isolation
+# flash-attn is excluded from uv's project resolution (pyproject.toml [tool.uv]
+# exclude-dependencies) -- but that exclusion ALSO silently suppresses an
+# explicit `uv pip install flash-attn` when run from the project dir. Install it
+# from /tmp so uv does not read the project's [tool.uv] and actually builds +
+# installs FA2 from source (ABI-matched to the installed torch; builds
+# sm_80/90/100/120 -> Ada/Hopper/Blackwell). Required by the Megatron
+# TransformerEngine `flash` attention backend (THD/packed sequences) and by
+# dense-Qwen3 flash_attention_2.
+RUN cd /tmp && uv pip install "flash-attn==2.8.3" --no-build-isolation
 
 CMD ["/bin/bash"]

docker/Dockerfile.sglang.megatron

@@ -19,6 +19,24 @@ ENV CUDA_HOME=/usr/local/cuda \
     NVTE_FRAMEWORK=pytorch \
     TORCH_CUDA_ARCH_LIST="8.0;8.9;9.0"
 
+# transformers 5.8 moved `rope_theta` into a nested `config.rope_parameters`
+# dict; mbridge 0.13.0 still reads the old top-level `config.rope_theta` and
+# crashes converting Qwen3 configs to Megatron (AttributeError) at model build.
+# Patch the installed mbridge to read either location. No published mbridge
+# release targets transformers 5.8 yet (AReaL avoids this by staying on
+# transformers <=5.3). mbridge is Megatron-only, so this lives in the Megatron
+# image rather than the FSDP base.
+RUN python - <<'PY'
+p = "/opt/venv/lib/python3.12/site-packages/mbridge/core/llm_bridge.py"
+s = open(p).read()
+old = "rotary_base=self.hf_config.rope_theta,"
+new = ('rotary_base=(getattr(self.hf_config, "rope_theta", None) '
+       'or getattr(self.hf_config, "rope_parameters", {}).get("rope_theta")),')
+assert old in s, "mbridge rope_theta line not found -- patch needs updating"
+open(p, "w").write(s.replace(old, new))
+print("[mbridge] patched rope_theta -> rope_parameters fallback")
+PY
+
 # Transformer Engine. The prebuilt `transformer-engine[pytorch]` wheels link
 # libcublas.so.12 and fail to load on a CUDA 13 base (ImportError:
 # libcublas.so.12), so build TE from source against the image's CUDA 13

docker/README.md

@@ -28,17 +28,19 @@ Pick the one that matches your **training backend**:
 
 ```bash
 # FSDP backend (default) — astraflow + SGLang + flash-attn. Covers most recipes.
-docker pull astraflowai/astraflow:v0.1.1
+docker pull astraflowai/astraflow:v0.1.2
 
 # Megatron-LM backend — the above plus Transformer Engine + apex.
 # Only needed when training with `backend: megatron` (TP/PP/EP, MoE, large models).
-docker pull astraflowai/astraflow:v0.1.1.megatron
+docker pull astraflowai/astraflow:v0.1.2.megatron
 ```
 
-`v0.1.1` is built from `Dockerfile.sglang`; `v0.1.1.megatron` from
+`v0.1.2` is built from `Dockerfile.sglang`; `v0.1.2.megatron` from
 `Dockerfile.sglang.megatron`. The Megatron image is a strict superset, so if you are
 unsure it also runs every FSDP recipe. Pin a version tag for reproducibility;
-`:latest` tracks the most recent FSDP release.
+`:latest` tracks the most recent FSDP release. Both `v0.1.2` images are validated
+end-to-end on 8×H100 (400-step math-RL runs incl. eval; FSDP with Qwen3.5-4B and
+dense Qwen3-8B, Megatron with Qwen3-8B).
 
 ## Build from source
 
@@ -55,12 +57,36 @@ docker build -f docker/Dockerfile.sglang.megatron -t astraflow:sglang-megatron .
 
 ## Quick Start
 
+The recommended workflow: **the image provides the environment** (Python venv,
+CUDA 13 toolkit, SGLang, flash-attn, fla kernels); **your local checkout provides
+the code**. astraflow is installed *editable* from `/workspace/astraflow`, so
+mounting your repo over that path makes the container run your code — code
+changes take effect immediately, and you only rebuild the image when the
+*environment* changes (dependency pins, CUDA, system libs).
+
+```bash
+# Train with YOUR local checkout inside the pre-built environment
+docker run --gpus all --net=host --shm-size=512g --ulimit nofile=65536:65536 \
+  -v /path/to/astraflow:/workspace/astraflow \
+  -v ~/.cache/huggingface:/hf -e HF_HOME=/hf \
+  -e WANDB_API_KEY=<your-key> \
+  astraflowai/astraflow:v0.1.2 \
+  bash examples/math/qwen3-8b-m2po-full/scripts/run_qwen3-8b-m2po-full.sh
+```
+
+- `-v /path/to/astraflow:/workspace/astraflow` — your repo replaces the baked-in
+  code (outputs land in `data-experiments/`/`data-log/` inside your checkout).
+- `-v ~/.cache/huggingface:/hf -e HF_HOME=/hf` — reuse your host model/dataset
+  cache instead of re-downloading inside the container.
+
+To poke around the image with its baked-in code instead (no mounts), start an
+interactive shell:
+
 ```bash
-# Run the pre-built FSDP image with host network and all GPUs
-docker run --gpus all --net=host --shm-size=512g --ulimit nofile=65536:65536 -it astraflowai/astraflow:v0.1.1
+docker run --gpus all --net=host --shm-size=512g --ulimit nofile=65536:65536 -it astraflowai/astraflow:v0.1.2
 
 # ...or the Megatron-backend image
-docker run --gpus all --net=host --shm-size=512g --ulimit nofile=65536:65536 -it astraflowai/astraflow:v0.1.1.megatron
+docker run --gpus all --net=host --shm-size=512g --ulimit nofile=65536:65536 -it astraflowai/astraflow:v0.1.2.megatron
 
 # ...or run a locally built image
 docker run --gpus all --net=host --shm-size=512g --ulimit nofile=65536:65536 -it astraflow:sglang

docs/en/get-started/installation.md

@@ -28,7 +28,7 @@ pip install -U "uv>=0.10"
 > (`extra-build-dependencies`, `override-dependencies`) that older uv
 > releases don't recognize. When uv hits an unknown `[tool.uv]` key it
 > silently ignores the *entire* `[tool.uv]` table, so the
-> `transformers==5.6.1` override (which must beat sglang's `==5.6.0` pin)
+> `transformers==5.8.1` override (the version the pinned sglang release requires)
 > is dropped and the install fails with an unsolvable
 > `transformers` conflict. The Docker images install the latest uv via the
 > official installer and are unaffected.
@@ -40,7 +40,7 @@ uv pip install -e ".[dev]"
 ```
 
 This installs all core dependencies (~260 packages) including PyTorch 2.11.0,
-Transformers 5.6.1, Megatron-Core 0.13.1, Ray, W&B, and dev tools (pytest, ruff,
+Transformers 5.8.1, Megatron-Core 0.13.1, Ray, W&B, and dev tools (pytest, ruff,
 ipython).
 
 ### Step 4: Install Flash Attention and SGLang
@@ -72,7 +72,7 @@ uv pip install "flash-attn==2.8.3" --no-build-isolation
 #### SGLang (inference backend)
 
 Install via the project extra so uv applies the `[tool.uv]` overrides (the
-`transformers==5.6.1` pin and the `flash-attn-4` pre-release allowance). SGLang
+`transformers==5.8.1` pin and the `flash-attn-4` pre-release allowance). SGLang
 pulls in FlashAttention-**4** (`flash-attn-4`, a pre-release wheel) automatically
 for its own attention backend — you do not install that one yourself.
 
@@ -87,7 +87,7 @@ pipeline / expert parallelism, MoE models). The default **FSDP** backend and
 all inference need nothing here — skip to Step 6.
 
 > **Prefer Docker?** Skip this entire step with the pre-built
-> `astraflowai/astraflow:v0.1.1.megatron` image (see Option B below), which already
+> `astraflowai/astraflow:v0.1.2.megatron` image (see Option B below), which already
 > bundles Transformer Engine + apex.
 
 `megatron-core` and `mbridge` are already installed by Step 3. The Megatron
@@ -167,15 +167,30 @@ Pre-built images are published on Docker Hub — they skip the from-source steps
 entirely. Requires the NVIDIA Container Toolkit so `--gpus all` works. Choose the image
 by **training backend**:
 
+The recommended workflow: **the image provides the environment, your local
+checkout provides the code**. astraflow is installed *editable* from
+`/workspace/astraflow`, so mounting your repo over that path runs your local
+code — no rebuild needed for code changes; rebuild only when dependency pins or
+the CUDA stack change.
+
 ```bash
-# FSDP backend (default) — covers most recipes
-docker run --gpus all --net=host --shm-size=512g --ulimit nofile=65536:65536 -it astraflowai/astraflow:v0.1.1
+# FSDP backend (default) — run a recipe from YOUR local checkout
+docker run --gpus all --net=host --shm-size=512g --ulimit nofile=65536:65536 \
+  -v /path/to/astraflow:/workspace/astraflow \
+  -v ~/.cache/huggingface:/hf -e HF_HOME=/hf \
+  -e WANDB_API_KEY=<your-key> \
+  astraflowai/astraflow:v0.1.2 \
+  bash examples/math/qwen3-8b-m2po-full/scripts/run_qwen3-8b-m2po-full.sh
 
 # Megatron-LM backend — adds Transformer Engine + apex (Step 5 above, pre-built in).
 # Use this for `backend: megatron` (TP/PP/EP, MoE, large models).
-docker run --gpus all --net=host --shm-size=512g --ulimit nofile=65536:65536 -it astraflowai/astraflow:v0.1.1.megatron
+docker run --gpus all --net=host --shm-size=512g --ulimit nofile=65536:65536 -it astraflowai/astraflow:v0.1.2.megatron
 ```
 
+Drop the `-v /path/to/astraflow:...` mount to use the code baked into the image
+instead; add `-it` for an interactive shell. The `-v ~/.cache/huggingface` mount
+reuses your host model/dataset cache instead of re-downloading in the container.
+
 > **Note on `--shm-size`:** this sets the size of the container's `/dev/shm`. A
 > recipe run co-locates the trainer, RaaS, and SGLang in a single container, all
 > sharing one `/dev/shm` — in particular RaaS stages received model weights under
@@ -191,6 +206,6 @@ docker run --gpus all --net=host --shm-size=512g --ulimit nofile=65536:65536 -it
 > `nofile` soft limit (1024) is far too low and the reward pool fails with
 > `[Errno 24] Too many open files`. Raise it with `--ulimit nofile=65536:65536`.
 
-The image bundles astraflow, SGLang, and flash-attn. Pin a version tag (`v0.1.1`) for
+The image bundles astraflow, SGLang, and flash-attn. Pin a version tag (`v0.1.2`) for
 reproducibility; `:latest` tracks the most recent release. See `docker/README.md` for
 build details and the NVIDIA Container Toolkit install guide.

docs/en/recipes/math.md

@@ -43,6 +43,47 @@ bash examples/math/qwen3-1.7b-m2po-2gpus-full/scripts/run_qwen3-1.7b-m2po-2gpus-
 | Train dataset | DeepScaleR |
 | Eval datasets | AIME24, AIME25, AMC, Minerva Math, MATH500 |
 
+## Qwen3.5-4B — 8 GPUs (hybrid Gated-DeltaNet)
+
+Math RL on **Qwen3.5-4B**, a hybrid Gated-DeltaNet + attention multimodal
+checkpoint trained text-only here. Requires the transformers 5.8 stack
+(`transformers==5.8.1`, `sglang==0.5.13.post1`, `fla` kernels — all pinned by
+`pyproject.toml`, or use the `astraflowai/astraflow:v0.1.2` image). Full and
+delta transfer variants:
+
+- [`qwen3.5-4b-m2po-full/`](https://github.com/Infini-AI-Lab/astraflow/tree/main/examples/math/qwen3.5-4b-m2po-full) — full weight transfer
+- [`qwen3.5-4b-m2po-delta/`](https://github.com/Infini-AI-Lab/astraflow/tree/main/examples/math/qwen3.5-4b-m2po-delta) — delta weight transfer
+
+### Run
+
+```bash
+bash examples/math/qwen3.5-4b-m2po-full/scripts/run_qwen3.5-4b-m2po-full.sh
+```
+
+### Settings
+
+| Setting | Value |
+|---|---|
+| Model | Qwen3.5-4B (`model_type: qwen3_5`, GDN hybrid) |
+| GPUs | 8 — RaaS ×4 (SGLang, DP=4), Trainer ×4 (FSDP, DP=4) |
+| Backend | **FSDP only** — Qwen3.5 is not supported on the Megatron backend (no `mbridge` bridge / GDN layer spec) |
+| Algorithm | M2PO (`m2_threshold` 0.01) |
+| Weight transfer | TCP — full or delta |
+| Context length | 8192 |
+| Max new tokens | 4000 |
+| Rollouts per prompt | 8 (`temperature` 1.0) |
+| Train batch size | 256 |
+| Learning rate | 5e-6 (Adam, constant schedule) |
+| Train steps | 800 |
+| Workflow / reward | `rlvr` / `math_verify` |
+| Train dataset | DeepScaleR |
+| Eval datasets | AIME24, AIME25, AMC, Minerva Math, MATH500 |
+
+Validated on 8×L40 and 8×H100 (on Hopper the launch scripts auto-set the
+`FLA_TILELANG`/`CUDA_HOME` environment the GDN backward needs). See the
+[recipe README](https://github.com/Infini-AI-Lab/astraflow/tree/main/examples/math/qwen3.5-4b-m2po-full)
+for the validated stack and results.
+
 ## Qwen3-8B — 8 GPUs
 
 The full-scale recipe. It needs an 8-GPU node — 4 GPUs for inference, 4 for training — and also comes in full and delta transfer variants:

examples/_common/utils.sh

@@ -108,3 +108,20 @@ astraflow_cleanup_trap() {
     wait 2>/dev/null
     exit 0
 }
+
+# Qwen3.5 (Gated-DeltaNet) on Hopper: the GDN backward must use fla's tilelang
+# kernel (fla blocks its triton path on sm_90 as numerically wrong, fla#640),
+# and the tilelang JIT needs a full CUDA toolkit (nvcc + CCCL headers; the
+# pip-shipped nvcc has none). Called by the Qwen3.5 recipe launchers before
+# starting the trainer. No-op on non-Hopper GPUs and when the user has already
+# set the variables.
+astraflow_setup_qwen35_hopper_env() {
+    if nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null | grep -q '^9\.'; then
+        export FLA_TILELANG="${FLA_TILELANG:-1}"
+        if [ -z "${CUDA_HOME:-}" ] && [ -d /usr/local/cuda ]; then
+            export CUDA_HOME=/usr/local/cuda
+            export PATH="${CUDA_HOME}/bin:${PATH}"
+        fi
+        echo "Hopper (sm_90) detected: FLA_TILELANG=${FLA_TILELANG} CUDA_HOME=${CUDA_HOME:-unset}"
+    fi
+}