From 9c0ebac5e15168794bea3f31b7abc5c7b23e142f Mon Sep 17 00:00:00 2001 From: zhouaihui Date: Sat, 7 Nov 2026 01:38:24 +0800 Subject: [PATCH] fix(recipes): use cuDNN for Qwen3.8-27B long-context training Select native cuDNN THD with CP2 and allow cold TP engine initialization. Validate observed attention against the saved configuration, rejecting backend mismatches while retaining full-context capacity requirements. --- aspec/commands/_recipe_capacity.py | 11 +++++++-- ...=> qwen3_8_27b_256k_cudnn_cp2_nvlink.yaml} | 13 +++++++---- docs/contributing/validation.md | 8 +++---- docs/user-guide/recipes.md | 5 ++-- tests/test_recipe_capacity.py | 23 +++++++++++++++---- tools/recipe_catalog.json | 2 +- 6 files changed, 43 insertions(+), 19 deletions(-) rename configs/recipes/{qwen3_8_27b_256k_fa4_cp2_nvlink.yaml => qwen3_8_27b_256k_cudnn_cp2_nvlink.yaml} (88%) diff --git a/aspec/commands/_recipe_capacity.py b/aspec/commands/_recipe_capacity.py index 149ad05..b6ee0ae 100644 --- a/aspec/commands/_recipe_capacity.py +++ b/aspec/commands/_recipe_capacity.py @@ -19,8 +19,14 @@ import json from pathlib import Path +from omegaconf import OmegaConf + def audit_long_run(directory): + config = OmegaConf.load(Path(directory) / "config.yaml") + backend = OmegaConf.select(config, "training.attention_backend") + if backend not in {"fa4", "cudnn_flex_attention"}: + raise ValueError(f"Unsupported configured attention backend for long-context CP2: {backend!r}") evidence = json.loads((Path(directory) / "evidence.json").read_text()) training = evidence.get("observations", {}).get("training", {}) scheduler = training.get("scheduler", {}) @@ -32,14 +38,15 @@ def audit_long_run(directory): or scheduler.get("rows_dispatched", 0) < 1 or scheduler.get("packed_tokens") != 261888 * scheduler.get("rows_dispatched", 0) or execution.get("context_parallel_size") != 2 - or execution.get("attention_backend") != "fa4" + or execution.get("attention_backend") != backend ): - raise ValueError("Missing actual 261888-token packed work with FA4 CP2") + raise ValueError(f"Missing actual 261888-token packed work with {backend} CP2") return { "tokens_per_row": 261888, "rows": scheduler["rows_dispatched"], "packed_tokens": scheduler["packed_tokens"], "context_parallel_size": 2, + "attention_backend": backend, } diff --git a/configs/recipes/qwen3_8_27b_256k_fa4_cp2_nvlink.yaml b/configs/recipes/qwen3_8_27b_256k_cudnn_cp2_nvlink.yaml similarity index 88% rename from configs/recipes/qwen3_8_27b_256k_fa4_cp2_nvlink.yaml rename to configs/recipes/qwen3_8_27b_256k_cudnn_cp2_nvlink.yaml index eb45d59..192590d 100644 --- a/configs/recipes/qwen3_8_27b_256k_fa4_cp2_nvlink.yaml +++ b/configs/recipes/qwen3_8_27b_256k_cudnn_cp2_nvlink.yaml @@ -1,9 +1,10 @@ # Release training starting point: three passes, held-out eval and checkpoints. # Hyperparameters are not a convergence/quality claim; see docs/user-guide/recipes.md. # Portable execution template. See configs/README.md for required inputs and scope. -# Adapted from AngelSpec's validated H200 Qwen3.8-27B 256K capacity-unlock -# lane and expressed through ASpec's canonical schema. One 8-GPU node: +# Adapted from AngelSpec's Qwen3.8-27B long-context configuration using +# ASpec's canonical schema. One 8-GPU node: # three TP2 Target engines (six GPUs) and one two-rank Direct CP trainer. +# Uses native cuDNN THD attention for the dense DSpark Draft. # Native prefill with optimized direct prefill disabled: one concurrent TP2 request # per replica and two source slots for the 261888-token workload. @@ -27,7 +28,7 @@ dataset: training: dflash_packing: true - attention_backend: fa4 + attention_backend: cudnn_flex_attention draft_context_parallel_size: 2 micro_batch_size: 1 draft_accumulation_steps: 2 @@ -73,6 +74,8 @@ inference: inference_batch_size: 1 sglang: tp_size: 2 + # Allow cold compilation during TP engine startup. + init_timeout: 900 mem_fraction_static: 0.90 extra_args: max_total_tokens: 262144 @@ -112,5 +115,5 @@ debug: performance_drain_steps: 1 timeline_enabled: false -output_dir: outputs/qwen3_8_27b_256k_fa4_cp2_nvlink -cache_dir: cache/qwen3_8_27b_256k_fa4_cp2_nvlink +output_dir: outputs/qwen3_8_27b_256k_cudnn_cp2_nvlink +cache_dir: cache/qwen3_8_27b_256k_cudnn_cp2_nvlink diff --git a/docs/contributing/validation.md b/docs/contributing/validation.md index 0d1b215..8757e7e 100644 --- a/docs/contributing/validation.md +++ b/docs/contributing/validation.md @@ -89,17 +89,15 @@ bash tools/run_in_env.sh python -m tools.recipes export \ The 27B bounded acceptance check (`accept qwen27b-long`) rejects a training Parquet containing any row shorter than 261,888 tokens. It also requires observed 261,888-token sequences and packed rows, positive dispatched work, and actual -FA4/CP2 execution in both phases. It records device memory every second in `gpu-memory.csv` and +CP2 execution with the configured attention backend in both phases (cuDNN by default). +A different observed backend fails the check. +It records device memory every second in `gpu-memory.csv` and per-device peaks/capacities in `acceptance.json`; missing observations fail acceptance. Sampled device-wide memory is distinct from an allocator peak and can include other processes, so use exclusive allocated GPUs. Synthetic inputs can establish capacity, but cannot establish quality. Ordinary `train qwen27b-long` runs may use shorter valid samples; the minimum length is an acceptance requirement, not a training input constraint. -A bounded 261,888-token train/resume/export check on eight SM90 GPUs with 143,771 MiB each -observed peaks of 118,578 MiB on Target and 99,987 MiB on Draft. These are observations for that -configuration, not minimum-memory requirements or a guarantee for a different workload. - To audit an existing train/resume/export receipt and its saved memory samples: ```bash diff --git a/docs/user-guide/recipes.md b/docs/user-guide/recipes.md index 7906d0d..d7585dc 100644 --- a/docs/user-guide/recipes.md +++ b/docs/user-guide/recipes.md @@ -27,11 +27,12 @@ connections on your hardware before training. | `qwen27b-long` | Qwen3.8-27B / DSpark | 1 node, 8 GPUs: 3 × TP2 Target + 2D (CP2); 256K-class | Distribute long-context training with CP | | `multinode` | Qwen3-8B / DSpark | 2 nodes × 8 GPUs; each 4T + 4D; 4K | Keep features node-local and synchronize gradients across nodes | -The 27B recipe uses a 261,888-token cap, FA4 and CP2. Each TP2 Target replica admits one request, +The [27B recipe](../../configs/recipes/qwen3_8_27b_256k_cudnn_cp2_nvlink.yaml) +uses a 261,888-token cap, native cuDNN THD attention and CP2. Each TP2 Target replica admits one request, with two source slots. `inference.target_prefill_memory_optimized` stays disabled because the direct prefill path currently requires TP1. Size feature buffers and backend token capacity for your inputs; see [long-context validation](../contributing/validation.md#algorithm-capacity-and-fault-checks) for -measured capacity and verification commands. Additional examples are under +capacity checks and verification commands. Additional examples are under [configs/examples](../../configs/examples/README.md). For a short single-node DSpark run with cuDNN attention, CP1 and no packing, use the diff --git a/tests/test_recipe_capacity.py b/tests/test_recipe_capacity.py index 617d28e..19b090e 100644 --- a/tests/test_recipe_capacity.py +++ b/tests/test_recipe_capacity.py @@ -19,8 +19,10 @@ from tools.validation.recipe_capacity import audit_long_run, summarize_memory -@pytest.mark.parametrize("broken", [None, "short", "empty", "cp1", "padding"]) -def test_long_capacity_requires_observed_full_context(tmp_path, broken): +@pytest.mark.parametrize("backend", ["fa4", "cudnn_flex_attention"]) +@pytest.mark.parametrize("broken", [None, "short", "empty", "cp1", "padding", "backend"]) +def test_long_capacity_requires_observed_full_context(tmp_path, broken, backend): + (tmp_path / "config.yaml").write_text(f"training:\n attention_backend: {backend}\n") training = { "scheduler": { "observed_sequence_min": 261888, @@ -29,7 +31,7 @@ def test_long_capacity_requires_observed_full_context(tmp_path, broken): "rows_dispatched": 4, "packed_tokens": 4 * 261888, }, - "execution": {"context_parallel_size": 2, "attention_backend": "fa4"}, + "execution": {"context_parallel_size": 2, "attention_backend": backend}, } if broken == "short": training["scheduler"]["observed_sequence_min"] = 128 @@ -39,12 +41,25 @@ def test_long_capacity_requires_observed_full_context(tmp_path, broken): training["execution"]["context_parallel_size"] = 1 elif broken == "padding": training["scheduler"]["packed_tokens"] -= 1 + elif broken == "backend": + training["execution"]["attention_backend"] = ( + "fa4" if backend == "cudnn_flex_attention" else "cudnn_flex_attention" + ) (tmp_path / "evidence.json").write_text(json.dumps({"observations": {"training": training}})) if broken: with pytest.raises(ValueError, match="actual"): audit_long_run(tmp_path) else: - assert audit_long_run(tmp_path)["tokens_per_row"] == 261888 + result = audit_long_run(tmp_path) + assert result["tokens_per_row"] == 261888 + assert result["attention_backend"] == backend + + +def test_long_capacity_rejects_unsupported_configured_backend(tmp_path): + (tmp_path / "config.yaml").write_text("training:\n attention_backend: math\n") + (tmp_path / "evidence.json").write_text("{}") + with pytest.raises(ValueError, match="configured attention backend"): + audit_long_run(tmp_path) def test_sampled_memory_requires_all_devices_and_repeated_samples(tmp_path): diff --git a/tools/recipe_catalog.json b/tools/recipe_catalog.json index 7aeb30f..98e3385 100644 --- a/tools/recipe_catalog.json +++ b/tools/recipe_catalog.json @@ -115,7 +115,7 @@ }, "qwen27b-long": { "configs": [ - "configs/recipes/qwen3_8_27b_256k_fa4_cp2_nvlink.yaml" + "configs/recipes/qwen3_8_27b_256k_cudnn_cp2_nvlink.yaml" ], "model": "Qwen3.8-27B", "nodes": 1,