diff --git a/MANIFEST.in b/MANIFEST.in index 295a3d6..527355c 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -15,7 +15,7 @@ include pyproject.toml recursive-include .github *.yml *.md recursive-include configs *.yaml *.md recursive-include docs *.md *.svg *.png *.json -recursive-include examples *.jsonl *.sh +recursive-include examples *.jsonl *.sh *.md recursive-include aspec/config *.json recursive-include aspec/compat/patches *.patch recursive-include licenses * diff --git a/README.md b/README.md index 1509ad0..76d0c88 100644 --- a/README.md +++ b/README.md @@ -59,6 +59,8 @@ Check [supported paths and limitations](docs/support.md), then browse the [recip 3. [Choose a training recipe](docs/user-guide/recipes.md) and [prepare training data](docs/user-guide/data.md). Recipes cover Qwen3-8B, Qwen3.5-9B Native MTP, Qwen3.8-27B long context, and Ling tiny/flash 3.0/flash 3.1. +For complete feature workflows, see [offline feature generation and replay](examples/qwen3-8b-offline/README.md) +and [DSA dense → indexer → sparse training](examples/qwen3-8b-dsa/README.md). From the repository root: diff --git a/aspec/offline/generate.py b/aspec/offline/generate.py index f281996..950c26c 100644 --- a/aspec/offline/generate.py +++ b/aspec/offline/generate.py @@ -143,7 +143,8 @@ def _cleanup_materialization( # Only release the reservation after source teardown has acknowledged. if inference_pg is not None and cleanup.get("state") == "completed": try: - ray.util.remove_placement_group(inference_pg) + pg, _, _ = inference_pg + ray.util.remove_placement_group(pg) except Exception as exc: errors.append(f"placement_cleanup: {type(exc).__name__}: {exc}") if errors: diff --git a/configs/examples/README.md b/configs/examples/README.md index 9a96298..3cae5ce 100644 --- a/configs/examples/README.md +++ b/configs/examples/README.md @@ -21,6 +21,11 @@ or disable evaluation/checkpointing; inspect each file before use. | [Qwen role-separated TCP](qwen3_8b_4k_multinode_tcp.yaml) | Separate Target and Draft nodes with shared storage | | [Qwen Target-generated quality protocol](qwen3_8b_dspark_target_generated.yaml) | Ten-pass training with held-out Target-generated data | | [TorchSpec-style input](sglang_qwen3_8b_dspark.yaml) | Fixed compatibility input, not rolling upstream support | +| [Offline materialization](qwen3_8b_offline_materialize.yaml) / [replay](qwen3_8b_offline_replay.yaml) | Generate DSpark features once, then train with zero Target GPUs | + +For complete source-checkout workflows, follow the +[offline generation, training and recovery example](../../examples/qwen3-8b-offline/README.md) +or the [DSA three-stage script](../../examples/qwen3-8b-dsa/README.md). The 9B and 27B DSpark examples each start with one eight-GPU node: four TP1 Target engines and four Draft GPUs, cuDNN attention, MBS1, CP1 and 4K context. Set `ASPEC_TARGET_MODEL` to the matching local checkpoint diff --git a/configs/examples/qwen3_8b_offline_materialize.yaml b/configs/examples/qwen3_8b_offline_materialize.yaml new file mode 100644 index 0000000..4132b90 --- /dev/null +++ b/configs/examples/qwen3_8b_offline_materialize.yaml @@ -0,0 +1,57 @@ +# Qwen3-8B DSpark feature materialization; see examples/qwen3-8b-offline/README.md. +# Uses the retained TorchSpec-derived offline generator and ASpec lifecycle validation. +# One Target GPU is used during generation; no Draft trainer is started. +model: + target_model_path: ${oc.env:ASPEC_TARGET_MODEL} + native_context_length: 32768 + trust_remote_code: false + draft_model_config: aspec/config/dspark_draft_config.json +dataset: + train_data_path: ${oc.env:ASPEC_TRAIN_DATA} + eval_data_path: ${oc.env:ASPEC_EVAL_DATA} + chat_template: qwen3-instruct + prompt_key: conversations + min_loss_tokens: 14 + shuffle_dataset: false +training: + attention_backend: cudnn_flex_attention + draft_context_parallel_size: 1 + dflash_packing: false + micro_batch_size: 1 + draft_accumulation_steps: 1 + training_num_gpus_per_node: 1 + training_num_nodes: 1 + fsdp_strategy: REPLICATE + max_seq_length: 4096 + num_epochs: 1 + dspark_num_target_layers: 5 + dspark_num_anchors: 64 + dflash_block_size: 7 +inference: + inference_engine_type: sgl + inference_num_gpus: 1 + inference_num_gpus_per_engine: 1 + inference_num_gpus_per_node: 1 + target_prefill_memory_optimized: false + store_last_hidden_states: true + inference_batch_size: 1 + inference_fetch_batch: 2 + inference_buffer_threshold: 2 + max_sample_pool_size: 4 + sglang: + tp_size: 1 + mem_fraction_static: 0.5 +mooncake: + data_plane: store + protocol: tcp + enable_gpu_direct: false + require_nvlink: false + hidden_dim: 4096 + max_seq_len: 4096 + global_segment_size: 8GB + local_buffer_size: 2GB + host_buffer_size: 2GB +logging: + report_to: none +output_dir: outputs/qwen3-8b-materialize +cache_dir: cache/qwen3-8b-materialize diff --git a/configs/examples/qwen3_8b_offline_replay.yaml b/configs/examples/qwen3_8b_offline_replay.yaml new file mode 100644 index 0000000..37a3e28 --- /dev/null +++ b/configs/examples/qwen3_8b_offline_replay.yaml @@ -0,0 +1,67 @@ +# Single-GPU DSpark replay; see examples/qwen3-8b-offline/README.md. +# Keep the Target checkpoint and capture settings identical to materialization. +model: + target_model_path: ${oc.env:ASPEC_TARGET_MODEL} + native_context_length: 32768 + trust_remote_code: false + draft_model_config: aspec/config/dspark_draft_config.json +dataset: + train_data_path: null + eval_data_path: null + renderer: null + min_loss_tokens: 14 + eval_interval: 4 + eval_micro_batch_size: 1 + shuffle_dataset: false +training: + attention_backend: cudnn_flex_attention + draft_context_parallel_size: 1 + dflash_packing: false + micro_batch_size: 1 + draft_accumulation_steps: 1 + training_num_gpus_per_node: 1 + training_num_nodes: 1 + fsdp_strategy: REPLICATE + max_seq_length: 4096 + num_epochs: 100 + num_train_steps: 8 + lr_total_steps: 8 + learning_rate: 0.0006 + min_lr: 0.00006 + warmup_ratio: 0.0 + seed: 42 + dspark_num_target_layers: 5 + dspark_num_anchors: 64 + dflash_block_size: 7 + prefetch_depth: 1 + streaming_dispatch_with_training: true + pipeline_dispatch_with_training: true + save_interval: 4 + save_final_checkpoint: true + save_per_epoch: false +inference: + inference_engine_type: offline + inference_num_gpus: 0 + target_prefill_memory_optimized: false + store_last_hidden_states: true + inference_batch_size: 1 + inference_fetch_batch: 2 + inference_buffer_threshold: 2 + max_sample_pool_size: 4 + offline: + data_path: ${oc.env:ASPEC_OFFLINE_DATA} + num_engines: 1 +mooncake: + data_plane: store + protocol: tcp + enable_gpu_direct: false + require_nvlink: false + hidden_dim: 4096 + max_seq_len: 4096 + global_segment_size: 8GB + local_buffer_size: 2GB + host_buffer_size: 2GB +logging: + report_to: none +output_dir: outputs/qwen3-8b-offline +cache_dir: cache/qwen3-8b-offline diff --git a/docs/user-guide/structures.md b/docs/user-guide/structures.md index 5c6c678..704c415 100644 --- a/docs/user-guide/structures.md +++ b/docs/user-guide/structures.md @@ -172,6 +172,10 @@ compressor/indexer, head, mapping and mixed dtype. See [export](evaluation.md). ## DSA: token-level sparse attention +The [three-stage runnable example](../../examples/qwen3-8b-dsa/README.md) connects +the existing dense, indexer and sparse configurations, with stage-local recovery +instructions and native export. + DSA uses a weighted-ReLU Lightning Indexer to select original context tokens **for each Draft query**. It neither compresses tokens like CSA/QSA nor shares one selection across a whole Draft block. The selected original K/V and the complete current bidirectional Draft block feed attention; diff --git a/docs/user-guide/training.md b/docs/user-guide/training.md index 48cb9aa..3ce581f 100644 --- a/docs/user-guide/training.md +++ b/docs/user-guide/training.md @@ -34,6 +34,9 @@ For a preplanned split, keep `training.lr_total_steps` equal to the **full** hor even when the first `num_train_steps` is smaller. Do not silently change data/objective/topology/model semantics across exact resume. Iteration names are one ahead of completed steps. Hidden `.incomplete` staging is not published recovery state. Uncommitted features may be recomputed. +With asynchronous packing, those samples may form different rows after restart; +checkpoint-state restoration does not guarantee the same document count or +numerical trajectory at a later optimizer step as an uninterrupted run. ## New phases and initial weights @@ -131,6 +134,10 @@ The GQA path executes `default`, `linear` and `yarn` RoPE from `rope_scaling` or ## Offline Target features +The [complete offline example](../../examples/qwen3-8b-offline/README.md) provides +paired materialization/replay configs and commands for evaluation, exact resume and +native export on one GPU used sequentially. + Materialize Target features once, then train Drafts from disk with zero Target GPUs. Draft architectures can share the dataset when they request the same Target feature contract. The current path uses single-node, single-Draft, CP1 training and CPU Store/TCP transport. Use a single-Draft SGLang config with local Target weights and train/eval data. Set `inference.target_prefill_memory_optimized=false`, `training.draft_context_parallel_size=1`, `training.dflash_packing=false`, and the `mooncake` settings shown below. Generate into a new or empty directory: diff --git a/examples/qwen3-8b-dsa/README.md b/examples/qwen3-8b-dsa/README.md new file mode 100644 index 0000000..99ddd48 --- /dev/null +++ b/examples/qwen3-8b-dsa/README.md @@ -0,0 +1,82 @@ +# Qwen3-8B DSA: three dependent training stages + +Run the existing full DSA Draft configuration through dense training, indexer +warmup and sparse training. Each stage uses a new optimizer/scheduler and output +directory, initializes from the preceding complete checkpoint, and retains source +progress. + +Use a source checkout with the [isolated Linux environment](../../docs/user-guide/installation.md) +activated. The default allocation is one Target GPU and one Draft GPU on a node +with NVLink, using FA4 packing at CP1. Supply disjoint train/eval inputs prepared +for Qwen3-8B, with sequences up to 4K and at least 32 supervised tokens per usable +sample. See [data preparation](../../docs/user-guide/data.md). + +```bash +export ASPEC_TARGET_MODEL=/path/to/Qwen3-8B +export ASPEC_TRAIN_DATA=/path/to/train.parquet +export ASPEC_EVAL_DATA=/path/to/eval.parquet +export CUDA_VISIBLE_DEVICES=0,1 +export RAY_ADDRESS=local + +bash examples/qwen3-8b-dsa/run.sh outputs/dsa-workflow +``` + +The script uses [the existing DSA YAML](../../configs/examples/qwen3_8b_dsa_dspark.yaml) +and three model JSONs. It checks each Plan, trains 16 updates per stage, evaluates +every four updates, saves checkpoints and exports the final sparse Draft as a native +asset. These short defaults demonstrate the workflow, not convergence or sparse +attention speed. Pass ordinary overrides after the output root, for example +`training.num_train_steps=8 training.lr_total_steps=8` for eight updates per stage. +The same overrides apply to all stages; the script fixes stage model, checkpoint, +source-cursor policy, capture requirements and output paths afterward. + +The output root must be new. Failed attempts remain available for diagnosis; the +script stops on the first failed command and never skips an existing stage. +The generated tree contains `dense/`, `indexer/`, `sparse/` and `native-asset/`. + +| Stage | Trainable state | Initialization | +| --- | --- | --- | +| Dense | Main Draft; indexer frozen | Fresh model | +| Indexer warmup | Indexer; main Draft frozen | Dense checkpoint, weights-only | +| Sparse | Main Draft and indexer | Indexer checkpoint, weights-only | + +Warmup does not request last hidden states needed by the normal DSpark objective; +the sparse stage requests them again. Warmup evaluation reports indexer KL rather than +a directly comparable language-model CE. Stage changes are intentional new phases, +not exact checkpoint recovery. + +## Resume within a stage + +If training was interrupted after a complete checkpoint, use that stage's saved +`config.yaml`. Clear **both** phase-transition flags for exact resume and select a +fresh output directory. Keep the original total step limit and LR horizon. + +Recovery restores checkpoint state and the committed source frontier. Asynchronous +packing can regroup uncommitted samples after restart, so a later optimizer-step +count can correspond to a different document count and numerical trajectory than +an uninterrupted run. Check restored state and source continuity separately from +later trajectory equivalence. + +```bash +aspec --config outputs/dsa-workflow/indexer/config.yaml \ + training.load_path=outputs/dsa-workflow/indexer/checkpoints \ + training.continual_training=false training.continue_source_cursor=false \ + output_dir=outputs/dsa-indexer-resumed +aspec --report outputs/dsa-indexer-resumed +``` + +After that stage finishes, start the sparse phase from its recovered checkpoint: + +```bash +aspec --config outputs/dsa-indexer-resumed/config.yaml \ + model.draft_model_config=aspec/config/dspark_qwen3_8b_dsa_sparse_draft_config.json \ + training.load_path=outputs/dsa-indexer-resumed/checkpoints \ + training.continual_training=true training.continue_source_cursor=true \ + inference.store_last_hidden_states=true output_dir=outputs/dsa-sparse-after-resume +``` + +This uses the same configured horizon for the new stage. Review each Report's +actual progress, required capabilities and lifecycle. DSA assets use +`aspec export --format aspec`; strict native loading is separate from downstream +serving integration. See [DSA semantics](../../docs/user-guide/structures.md#dsa-token-level-sparse-attention) +and [recovery](../../docs/user-guide/training.md). diff --git a/examples/qwen3-8b-dsa/run.sh b/examples/qwen3-8b-dsa/run.sh new file mode 100644 index 0000000..310537d --- /dev/null +++ b/examples/qwen3-8b-dsa/run.sh @@ -0,0 +1,59 @@ +#!/usr/bin/env bash +# Copyright 2026 Ant Group Co., Ltd. Licensed under Apache-2.0. +# Run dense, indexer warmup and sparse training with native checkpoint export. +set -euo pipefail + +if [[ $# -lt 1 || ${1:-} == --help || ${1:-} == -h ]]; then + echo 'Usage: bash examples/qwen3-8b-dsa/run.sh NEW_OUTPUT_ROOT [KEY=VALUE ...]' + echo 'Set ASPEC_TARGET_MODEL, ASPEC_TRAIN_DATA, ASPEC_EVAL_DATA and CUDA_VISIBLE_DEVICES.' + echo 'Runs dense -> indexer warmup -> sparse, then exports a native asset.' + exit 0 +fi +: "${ASPEC_TARGET_MODEL:?Set an absolute local Qwen3-8B checkpoint path}" +: "${ASPEC_TRAIN_DATA:?Set an absolute training data path}" +: "${ASPEC_EVAL_DATA:?Set an absolute disjoint evaluation data path}" +: "${CUDA_VISIBLE_DEVICES:?Select two allocated GPUs explicitly}" + +# Resolve the output before changing directory so paths are relative to the caller. +output_root=$(python -c 'import pathlib,sys; print(pathlib.Path(sys.argv[1]).resolve())' "$1") +shift +repo_root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." && pwd)" +cd "$repo_root" +mkdir -p -- "$(dirname -- "$output_root")" +mkdir -- "$output_root" +export RAY_ADDRESS="${RAY_ADDRESS:-local}" +config=configs/examples/qwen3_8b_dsa_dspark.yaml +previous= + +for stage in dense indexer sparse; do + continuation=false + last_hidden=true + if [[ $stage != dense ]]; then + continuation=true + fi + if [[ $stage == indexer ]]; then + last_hidden=false + fi + args=(--config "$config" + "model.target_model_path=$ASPEC_TARGET_MODEL" + "dataset.train_data_path=$ASPEC_TRAIN_DATA" + "dataset.eval_data_path=$ASPEC_EVAL_DATA" + dataset.eval_interval=4 "$@" + "model.draft_model_config=aspec/config/dspark_qwen3_8b_dsa_${stage}_draft_config.json" + "training.continual_training=$continuation" + "training.continue_source_cursor=$continuation" + "training.load_path=${previous:-null}" + training.save_final_checkpoint=true + "inference.store_last_hidden_states=$last_hidden" + "output_dir=$output_root/$stage" + "cache_dir=$output_root/cache") + aspec "${args[@]}" --dry-run + aspec "${args[@]}" + aspec --report "$output_root/$stage" + previous="$output_root/$stage/checkpoints" +done + +# Export the final published checkpoint selected by the existing tracker. +iteration=$(cat "$previous/latest_checkpointed_iteration.txt") +printf -v checkpoint '%s/iter_%07d' "$previous" "$iteration" +aspec export --format aspec --input-dir "$checkpoint" --output-dir "$output_root/native-asset" diff --git a/examples/qwen3-8b-offline/README.md b/examples/qwen3-8b-offline/README.md new file mode 100644 index 0000000..f8dbfe4 --- /dev/null +++ b/examples/qwen3-8b-offline/README.md @@ -0,0 +1,70 @@ +# Generate features once, train from disk + +This source-checkout example uses one GPU in two sequential phases: Qwen3-8B +materializes DSpark features, then one Draft trains without a live Target engine. +It uses single-node CP1, no packing, and CPU Store/TCP transport. The eight-update +training window demonstrates evaluation, exact recovery and native export; it is +not a quality or performance recipe. + +Activate the [isolated Linux environment](../../docs/user-guide/installation.md) +and run from the repository root. Set absolute paths to the **local** Qwen3-8B +checkpoint, disjoint train/eval conversations or prepared Parquet, and a new feature +directory. A small input subset is useful for a first run: generation saves **all** +admitted samples, regardless of the training step limit. Features can occupy much +more disk space than the source text. + +```bash +export ASPEC_TARGET_MODEL=/path/to/Qwen3-8B +export ASPEC_TRAIN_DATA=/path/to/train.jsonl +export ASPEC_EVAL_DATA=/path/to/eval.jsonl +export ASPEC_OFFLINE_DATA=/path/to/new-qwen3-features +export CUDA_VISIBLE_DEVICES=0 +export RAY_ADDRESS=local + +bash tools/run_in_env.sh python -m aspec.offline.generate \ + --config configs/examples/qwen3_8b_offline_materialize.yaml \ + --output "$ASPEC_OFFLINE_DATA" +``` + +Use the [documented data formats](../../docs/user-guide/data.md). Each usable +training sample needs at least 14 supervised tokens, and sequences must fit the +4K limit. `materialization.json` records saved train/eval counts and cleanup; +`dataset.json`, `manifest.jsonl` and `samples/` form the replay dataset. An interrupted +generation is incomplete: preserve it and generate into a fresh directory. + +## Train, then resume + +Inspect the Plan after generation completes. It should select offline replay, +zero Target GPUs and one Draft GPU. The same GPU used for generation can now train. + +```bash +aspec --config configs/examples/qwen3_8b_offline_replay.yaml --dry-run + +aspec --config configs/examples/qwen3_8b_offline_replay.yaml \ + training.num_train_steps=4 output_dir=outputs/offline-first +aspec --report outputs/offline-first + +aspec --config configs/examples/qwen3_8b_offline_replay.yaml \ + training.load_path=outputs/offline-first/checkpoints \ + training.continual_training=false training.continue_source_cursor=false \ + output_dir=outputs/offline-resumed +aspec --report outputs/offline-resumed + +aspec export --format aspec \ + --input-dir outputs/offline-resumed/checkpoints/iter_0000009 \ + --output-dir outputs/offline-native-asset +``` + +Both segments retain `lr_total_steps=8`; the resumed run continues to total step 8. +Checkpoint iteration names are one ahead of completed optimizer steps. Choose fresh +output/export directories on subsequent attempts. Evaluation reads the recorded +eval split rather than starting a Target engine. + +Keep the Target checkpoint available: replay still loads frozen embedding/head +weights. Target weights, tokenizer, capture layers, normalization, dtype and dataset +identity must match generation. Changing preprocessing requires new features. +Architectures with compatible feature contracts may reuse this dataset in separate +runs; changing the feature contract requires materializing it again. + +See [training](../../docs/user-guide/training.md#offline-target-features) and +[export boundaries](../../docs/user-guide/evaluation.md). diff --git a/tests/test_startup_cleanup.py b/tests/test_startup_cleanup.py index 53d4036..0111fb7 100644 --- a/tests/test_startup_cleanup.py +++ b/tests/test_startup_cleanup.py @@ -116,7 +116,7 @@ def fence(statuses, **kwargs): [engine], saver, owner, - pg, + (pg, [0], [0]), ) assert result["state"] == "completed_with_errors" if fence_ok: @@ -134,6 +134,32 @@ def fence(statuses, **kwargs): assert result["source_teardown"] == "blocked" +@pytest.mark.parametrize("source_state", ["completed", "completed_with_errors"]) +def test_materialization_releases_only_the_placement_handle_after_source_cleanup( + monkeypatch, source_state +): + from aspec.offline import generate as module + + pg = object() + events = [] + monkeypatch.setattr( + module, "_safe_training_cleanup", + lambda *a, **k: events.append("sources") or {"state": source_state}, + ) + remove = MagicMock(side_effect=lambda group: events.append("placement")) + monkeypatch.setattr(module.ray.util, "remove_placement_group", remove) + # create_placement_groups returns (handle, bundle indices, device ids). + result = module._cleanup_materialization( + SimpleNamespace(), None, None, None, [], None, None, (pg, [0], [2]) + ) + assert result["state"] == source_state + if source_state == "completed": + remove.assert_called_once_with(pg) + assert events == ["sources", "placement"] + else: + remove.assert_not_called() + + def test_main_disconnects_job_actors_after_initialization_error(monkeypatch, tmp_path): from aspec import train_entry as module from aspec.ray import node_local_files