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..401789c 100644 --- a/README.md +++ b/README.md @@ -25,7 +25,8 @@ independent experiments, and take a checkpoint through evaluation and export. - **One Target, multiple concurrent Draft experiments.** Compare algorithms, structures or hyperparameters while generating compatible Target features once. Each Draft keeps its own weights, optimizer, progress and checkpoints, without a shared optimizer-step barrier. - See [shared Target training](docs/user-guide/sharing.md) for the single-node workflow. + See [shared Target training](docs/user-guide/sharing.md) and the + [comparison and recovery walkthrough](examples/qwen3-8b-shared-target/README.md). - **Make long-context training practical.** Combine sequence packing, context parallelism and optimized attention execution to reduce padding waste and distribute long sequences across GPUs. - **Explore without rebuilding the workflow.** Switch between algorithm and model recipes with diff --git a/configs/recipes/README.md b/configs/recipes/README.md index 170e174..e9c3f3d 100644 --- a/configs/recipes/README.md +++ b/configs/recipes/README.md @@ -9,5 +9,7 @@ The optional validation command exports and reloads a native asset. `native_only whose supported export boundary is native-only; `false` does not disable the native roundtrip. Use each recipe's documented inputs and check capacity on your own hardware. +The [shared Target walkthrough](../../examples/qwen3-8b-shared-target/README.md) uses +the recipe for experiment comparison, separate curves and recovery from different checkpoints. Supplemental variants belong under `configs/examples/`; benchmark comparisons remain under `configs/benchmarks/`. diff --git a/docs/user-guide/sharing.md b/docs/user-guide/sharing.md index 86d1606..d0158b0 100644 --- a/docs/user-guide/sharing.md +++ b/docs/user-guide/sharing.md @@ -4,6 +4,10 @@ Shared mode captures compatible Target features once and trains independent Draf Each experiment owns its weights, optimizer, LR, RNG, evaluation, work frontier, checkpoint and tracking output. Shared execution is currently single-node. +The [comparison and independent recovery walkthrough](../../examples/qwen3-8b-shared-target/README.md) +derives two learning-rate experiments from the shared-target recipe, records separate +TensorBoard curves and resumes from unequal checkpoints while one Draft is already complete. + Use this workflow to compare Draft algorithms, structures or hyperparameters for the same Target. Algorithm selection chooses what each experiment trains; shared execution lets several such experiments run concurrently with compatible Target capture generated once. Each keeps its own @@ -60,7 +64,8 @@ Copy the original saved `config.yaml` outside the run directory and edit each ex list, model/data/ordering, LR horizon and execution semantics. Current dotlist overrides do not support patching an individual list index; replacing the list with partial entries can lose fields. -Set `training.continual_training=false`, keep the full `lr_total_steps`, choose a fresh root output +Set both `training.continual_training=false` and `training.continue_source_cursor=false` +for every child, keep the full `lr_total_steps`, choose a fresh root output directory, and set `num_train_steps` to the intended total frontier. Children may resume from different committed frontiers. One child's progress never licenses skipping another child's work. diff --git a/examples/qwen3-8b-shared-target/README.md b/examples/qwen3-8b-shared-target/README.md new file mode 100644 index 0000000..94e7163 --- /dev/null +++ b/examples/qwen3-8b-shared-target/README.md @@ -0,0 +1,182 @@ +# Compare Draft experiments and resume them independently + +This source-checkout walkthrough reuses the existing +[shared-target recipe](../../configs/recipes/qwen3_8b_4k_multi_draft_nvlink.yaml). +Two DSpark Drafts share compatible Qwen3-8B Target features, use different learning +rates, and own separate training state, evaluation curves and checkpoints. +The short run demonstrates experiment management; it does not establish which +learning rate converges better. + +Use the [isolated Linux environment](../../docs/user-guide/installation.md), with +the [TensorBoard dependency](../../docs/user-guide/training.md#tensorboard-and-wb) +installed, and run from the repository root. The recipe allocates eight GPUs on +one NVLink node: four Target GPUs and two GPUs for each Draft. It uses CP1, MBS1, +four accumulation steps and no packing, giving eight documents per optimizer +update per Draft. + +Provide a local Qwen3-8B checkpoint and disjoint train/eval conversation files +following the [data format guide](../../docs/user-guide/data.md). Samples must fit +4K context and contain at least 32 supervised tokens. For a first run, a small +subset such as 128 training and 16 evaluation samples is sufficient. + +```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_SHARED_OUTPUT=outputs/qwen3-shared-compare +export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 +export RAY_ADDRESS=local +``` + +## Prepare and run the comparison + +Create a new directory and derive a resolved configuration from the recipe. +Keep both complete experiment entries. This example gives both Drafts seed 42 and +a shared source-order seed, so the intended hyperparameter difference is the +learning-rate schedule: A uses `6e-4 → 6e-5`, B uses `3e-4 → 3e-5`. Both schedules +have the same eight-update horizon. B stops at update four to provide a different +recovery frontier. + +```bash +python - <<'PY' +import os +from pathlib import Path +from omegaconf import OmegaConf +from aspec.config.train_config import load_config + +root = Path(os.environ["ASPEC_SHARED_OUTPUT"]).expanduser().resolve() +root.mkdir(parents=True, exist_ok=False) +config = load_config( + "configs/recipes/qwen3_8b_4k_multi_draft_nvlink.yaml", + cli_args=[ + "dataset.chat_template=qwen3-instruct", + "dataset.dataset_seed=42", + "dataset.eval_interval=4", + "training.seed=42", + "training.num_train_steps=8", + "training.lr_total_steps=8", + "training.save_interval=4", + "training.save_per_epoch=false", + "logging.report_to=tensorboard", + f"output_dir={root / 'first'}", + f"cache_dir={root / 'cache'}", + ], + save_snapshot=False, +) +for experiment in config.experiments: + experiment.training.seed = 42 + experiment.training.num_train_steps = 8 if experiment.name == "draft_a" else 4 +OmegaConf.save(config, root / "first.yaml", resolve=True) +PY +``` + +Inspect the shared Plan and run it: + +```bash +aspec --config "$ASPEC_SHARED_OUTPUT/first.yaml" --dry-run +aspec --config "$ASPEC_SHARED_OUTPUT/first.yaml" +aspec --report "$ASPEC_SHARED_OUTPUT/first" +aspec --report "$ASPEC_SHARED_OUTPUT/first/draft_a" +aspec --report "$ASPEC_SHARED_OUTPUT/first/draft_b" +``` + +Expected optimizer progress is A `0 → 8` and B `0 → 4`. Each child directory has +its own `config.yaml`, `plan.json`, `evidence.json`, `tracking.json`, `runs/` and +`checkpoints/`. Verify healthy completion, evaluation and cleanup in the reports. +`tracking.json` records each writer's location and closure state. + +Open local curves with: + +```bash +tensorboard --logdir "$ASPEC_SHARED_OUTPUT/first" +``` + +Select `draft_a/runs` and `draft_b/runs` and compare the same metric at matching +optimizer steps, especially held-out evaluation at step four. In this configuration, +training scalar steps, evaluation and checkpoint frontiers count completed updates. +A's step-eight result cannot be compared with B's step-four +result as an equal-budget quality comparison. Common seeds do not establish +identical asynchronous execution or a statistically reliable quality ranking. +Shared Target production and per-Draft consumption have different scopes; do not +sum child consumption and call it unique Target work or isolated throughput. + +## Resume B while A is already complete + +Use the **saved root configuration**, which contains the full experiment list. +Point each child at its own explicit checkpoint: A at eight completed updates, +B at four. Checkpoint iteration names are one ahead of completed updates. +Preserve all model/data/order, execution, optimizer and LR-horizon settings. +Clear both phase-transition flags within every experiment. + +```bash +python - <<'PY' +import os +from pathlib import Path +from omegaconf import OmegaConf + +root = Path(os.environ["ASPEC_SHARED_OUTPUT"]).expanduser().resolve() +config = OmegaConf.load(root / "first/config.yaml") +frontiers = {"draft_a": 9, "draft_b": 5} +assert {experiment.name for experiment in config.experiments} == set(frontiers) +for experiment in config.experiments: + checkpoint = root / "first" / experiment.name / "checkpoints" / f"iter_{frontiers[experiment.name]:07d}" + if not checkpoint.is_dir(): + raise FileNotFoundError(checkpoint) + experiment.training.load_path = str(checkpoint) + experiment.training.num_train_steps = 8 + experiment.training.continual_training = False + experiment.training.continue_source_cursor = False +config.output_dir = str(root / "resumed") +with (root / "resume.yaml").open("x") as stream: + stream.write(OmegaConf.to_yaml(config, resolve=True)) +PY +``` + +```bash +aspec --config "$ASPEC_SHARED_OUTPUT/resume.yaml" --dry-run +aspec --config "$ASPEC_SHARED_OUTPUT/resume.yaml" +aspec --report "$ASPEC_SHARED_OUTPUT/resumed" +aspec --report "$ASPEC_SHARED_OUTPUT/resumed/draft_a" +aspec --report "$ASPEC_SHARED_OUTPUT/resumed/draft_b" +``` + +Expected progress is A `8 → 8` with `no_new_work`, and B `4 → 8`. A's committed +work must remain intact; B restores its own model, optimizer, scheduler, RNG and +source frontier. This demonstrates that finishing one experiment does not advance +another experiment's progress. The session retains the original eight-GPU +allocation and initializes both children; completed A is not a GPU-rescheduling +mechanism. A failed or cancelled child makes the shared session exit nonzero. + +A's report is `completed / incomplete_evidence` because it performs no new +training or feature transfer; the root report is therefore `incomplete_or_failed`. +Check A's `no_new_work`, unchanged committed frontier, closed tracking and completed +cleanup. B should report `completed / healthy`. + +Keep the full experiment list even when only B has new work. Individual list-index +dotlist overrides are unsupported, and replacing `experiments` with partial +entries can discard the original configuration. To resume both children from +different frontiers, choose checkpoints below their intended total horizons. +See [recovery semantics](../../docs/user-guide/training.md#exact-resume); this +workflow restores each child independently; it does not guarantee equality with +an uninterrupted training trajectory. + +TensorBoard records the resumed attempt separately, preserving restored metric +steps rather than appending to the old event files. To view both attempts, point +`tensorboard --logdir` at `$ASPEC_SHARED_OUTPUT`; B's training steps continue at +five and its next evaluation is at eight. Use A's first attempt for its completed +curves, since its resumed attempt performs no new optimizer updates. + +For optional native export, A's completed asset comes from the first run and B's +from the resumed run: + +```bash +aspec export --format aspec \ + --input-dir "$ASPEC_SHARED_OUTPUT/first/draft_a/checkpoints/iter_0000009" \ + --output-dir "$ASPEC_SHARED_OUTPUT/asset-a" +aspec export --format aspec \ + --input-dir "$ASPEC_SHARED_OUTPUT/resumed/draft_b/checkpoints/iter_0000009" \ + --output-dir "$ASPEC_SHARED_OUTPUT/asset-b" +``` + +Use a fresh output root for another attempt. See [shared Target support](../../docs/user-guide/sharing.md) +and [export boundaries](../../docs/user-guide/evaluation.md). diff --git a/tests/test_multi_draft_plan.py b/tests/test_multi_draft_plan.py index 56dd133..268a2c9 100644 --- a/tests/test_multi_draft_plan.py +++ b/tests/test_multi_draft_plan.py @@ -224,16 +224,18 @@ def test_shared_context_semantics_are_validated_before_gpu_allocation(): compile_multi_draft_plan(config) -def test_full_source_does_not_force_first_draft_to_use_fa4(tmp_path): +@pytest.mark.parametrize("attention_backend", ["flex_attention", "cudnn_flex_attention"]) +def test_shared_source_preserves_first_draft_attention_backend(tmp_path, attention_backend): import json root = Path(__file__).parents[1] raw = json.loads((root / "aspec/config/dflash2_draft_config.json").read_text()) - raw["is_causal"] = True # This existing DFlash2 combination requires Flex. + raw["is_causal"] = True draft = tmp_path / "causal.json" draft.write_text(json.dumps(raw)) config = load_config(str(root / "configs/examples/qwen3_8b_dflash2.yaml"), save_snapshot=False) config.model.draft_model_config = str(draft) + config.training.attention_backend = attention_backend config.inference.store_last_hidden_states = True config.training.training_num_gpus_per_node = 1 config.experiments = OmegaConf.create([ @@ -245,7 +247,7 @@ def test_full_source_does_not_force_first_draft_to_use_fa4(tmp_path): ]) plan = compile_multi_draft_plan(config) assert plan.to_runtime_args().draft_context_parallel_size == 1 - assert plan.to_runtime_args().attention_backend == "flex_attention" + assert plan.to_runtime_args().attention_backend == attention_backend assert plan.plans["cp2"].to_runtime_args().draft_context_parallel_size == 2