axolotl-ai-cloud/axolotl
1
0
Fork 0
mirror of https://github.com/axolotl-ai-cloud/axolotl.git synced 2026-10-12 02:11:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: diffusion_decision plugin for typed-decision training on native diffusion LMs (#4117)

Browse Source 
* feat(diffusion): add diffusion_decision plugin for typed-decision training

* feat(diffusion): add public procedural decision mix, eval script and decision recipes

* docs(diffusion): document typed-decision training

* docs(examples): explain why the decision recipes pin varlen attention

* test: decision trainer suppresses token perplexity; base diffusion trainer no longer reports it

* test(diffusion_decision): share spec, record, canvas, tokenizer and harness factories across files

* test(diffusion_decision): fold repeated setup into builders and parametrize near-identical cases

* refactor(diffusion_decision): drop the evaluation harness, HF reader and data audit

Training needs none of them; the prepared cache now validates on its identity
and row hashes alone.

* fix(diffusion_decision): address review on dev split, record ids and merge hook

The automatic family dev split can no longer hold out every record; grouping
synthesizes positional ids for rows without one and rejects duplicates;
post_lora_merge skips runs without a decision config. Types the slot
placement in the trainer test for mypy.

* fix: use Ray native Pydantic support on Python 3.14

(cherry picked from commit 9bdcd4921d889d680f81e04ef4abd581e5de1db4)

* feat(diffusion_decision): per-question label weights in the decision loss

* fix(diffusion_decision): address review on collator fallback, pinned labels and the question cap

The collator hook returns None without a decision block so defaults apply;
canvas validation rejects pinned label positions, which the collator would
supervise uncorrupted; max_questions_per_canvas is capped at the template's
64-question limit.
This commit is contained in:
Wing Lian
2026-10-06 13:41:51 -04:00
committed by GitHub
parent 4608acfd55
commit ddbc5e19f8
77 changed files with 16691 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
MANIFEST.in
+1
View File
@@ -5,6 +5,7 @@ include src/axolotl/core/trainers/prm/NOTICE
include VERSION
include src/axolotl/utils/chat_templates/templates/*.jinja
include src/axolotl/integrations/kernels/libs/scattermoe_lora/metadata.json
include src/axolotl/integrations/diffusion_decision/vendored/DJEV_LICENSE
recursive-include src/axolotl *.yaml
include AGENTS.md
recursive-include docs/agents *.md
docs/diffusion_lm.qmd
+170 -5
View File
@@ -200,8 +200,8 @@ lora_target_modules: '^model\.(encoder\.language_model|decoder)\.layers\.\d+\.se
Vision, router and expert parameters stay frozen.
Nemotron supports opt-in Cut Cross Entropy for its aligned diffusion objective.
With Axolotl's CCE fork installed, add:
Nemotron supports opt-in Cut Cross Entropy for its aligned diffusion objective
and for typed-decision training. With Axolotl's CCE fork installed, add:
```yaml
plugins:
@@ -211,9 +211,10 @@ cut_cross_entropy: true
Keep any other plugins the recipe requires in that list. The final training
read computes full-vocabulary CE from hidden states without materializing
vocabulary logits. Targets stay aligned, with no causal shift. This path
requires a plain linear output head and does not support a LoRA output head or
the `loo` objective.
vocabulary logits; decision training still computes its restricted-label and
Brier terms over the allowed answers. Targets stay aligned, with no causal
shift. This path requires a plain linear output head and does not support a
LoRA output head or the `loo` objective.
`axolotl merge-lora config.yaml` exports a merged checkpoint. When
`merge_method` is omitted, DiffusionGemma selects `legacy` so its tied factors
@@ -224,6 +225,170 @@ reply is a completion block appended to the conversation and denoised in one
piece; see [Inference](inference.qmd#sec-chat-diffusion). Native models
generate completions only.
## Typed-decision training
The `diffusion_decision` plugin trains a native diffusion model to answer
typed multiple-choice, score and yes/no questions. Each record renders a fixed
decision system prompt from its question schema, sends its `state` as the user
message, and places one answer label per question in the output canvas using
the canvas template of the public `mmastrac/djev` project.
`examples/nemotron-diffusion/decision-lora-8b.yaml` is the starting recipe for
`nvidia/Nemotron-Labs-Diffusion-8B`. It uses BF16, `attn_implementation: varlen`,
and rank-64/alpha-128 LoRA over all seven attention and MLP projections with
`loraplus_lr_ratio: 8`. It trains one epoch with AdamW at `6.5e-6`, microbatch
16 and eight accumulation steps (effective batch 128). The decision loss
combines candidate-restricted CE, full-vocabulary CE and a `0.1` Brier term,
with `unroll.k_max: 2`, no latent slots and non-reentrant gradient
checkpointing. The YAML leaves the model revision unpinned; add
`revision_of_model` to a local copy for an exact reproduction.
```bash
axolotl preprocess examples/nemotron-diffusion/decision-lora-8b.yaml
axolotl train examples/nemotron-diffusion/decision-lora-8b.yaml
```
Set the `data_files` of the `datasets` and `test_datasets` entries to separate
normalized train and development JSONL files. The recipe evaluates every 50
steps and keeps the adapter with the best `eval_loss`. Keep held-out test data
out of both entries.
### Dataset format
`type: diffusion_decision.jsonl` reads already-normalized JSONL, one record per
line. `diffusion_decision.typed_decisions` and `diffusion_decision.procedural`
normalize rows from the `typed_decisions` and procedural typed-decision source
schemas into the same form.
```json
{"source":"support_demo","family":"returns","id":"case-0042","group":"case-0042","instructions":"Decide from the supplied case only.","state":{"customer":"A customer says their unopened order arrived yesterday and asks for a refund.","policy":"Unopened items delivered within 30 days may be refunded."},"questions":{"refund":{"type":"choice","instructions":"Choose the next action.","options":[{"name":"approve","description":"Approve a refund under the stated policy."},{"name":"deny","description":"Deny the request because the policy does not apply."},{"name":"escalate","description":"Send the case to a human reviewer."}]},"urgency":{"type":"score","instructions":"Rate how urgently this needs review.","levels":["routine","soon","urgent"]},"eligible":{"type":"noul","instructions":"Is the customer eligible for a refund?","criteria":{"true":"The stated policy permits a refund.","false":"The stated policy does not permit a refund."}}},"labels":{"refund":{"kind":"dist","probs":[0.78,0.02,0.20]},"urgency":{"kind":"hard","gold_idx":0},"eligible":{"kind":"set","allowed_set":[0]}},"source_metadata":{"provenance":{"dataset":"my-org/support-decisions","revision":"immutable-revision","split":"train","source_id":"case-0042"}}}
```
`source`, `id` and `group` identify the row; `family` is optional and is used
for split isolation. Records with the same source, family, state and top-level
`instructions` that share a `group` are rendered on one canvas, up to
`max_questions_per_canvas` questions. There is no per-row `system` field: put
shared context in `state`, global instructions in `instructions`, and
per-question instructions on each question. Keep question IDs stable and
nonempty.
| Type | Required schema field | Canonical alternative order |
|---|---|---|
| `choice` | `options` | List order. An option is a string or an object with `name` and `description`; descriptions are rendered. |
| `score` | `levels` | List order, lowest to highest as defined by the source. Do not sort levels lexically. |
| `noul` | none | `yes`, then `no`. Optional `criteria` documents the true/false conditions and is rendered. |
Label indexes always refer to this canonical order, never to option names or
token IDs. Answer tokens come from `labels.codebook`: `vendored26` uses the
djev letters `A`-`Z` (up to 26 alternatives); `expanded52` allows up to 52.
### Labels
Every question has exactly one label:
```json
{"kind":"hard","gold_idx":1}
{"kind":"dist","probs":[0.10,0.70,0.20]}
{"kind":"set","allowed_set":[0,2]}
```
`hard` names one zero-based alternative. `set` is a nonempty set of unique
valid indexes. `dist` has one finite, nonnegative probability per alternative
that sums to one within `1e-4`; zero entries are valid and must not be omitted.
Partial distributions, ranking targets, per-row candidate token IDs and
hard/soft blending are not accepted.
`labels.label_softmax` selects candidate-restricted CE (`restricted`),
full-vocabulary CE over the allowed label tokens (`full`), or both (the
default). The full-vocabulary term still penalizes probability outside the
valid answers. `full_ce_weighting: dft` applies DFT weighting to hard-label
full CE. `hard_label_smoothing` spreads that fraction of a one-hot CE target
uniformly across the question's valid answer tokens; Brier targets are
unchanged.
### Splits and provenance
Keep train, development and test data in separate files or dataset entries,
and use `split: train` only for training entries. Preserve
`source_metadata.provenance` with an immutable source location and revision,
the split, and the source row ID. `typed_decisions` rows train only from a
declared `train` split and reject a row `split` that disagrees with it. After
normalization, the loader rejects states or grouped contexts shared between
splits.
### Mixture and weighting
`mixture.weights` sets per-source sampling weights by `source`. Without
weights, each source is sampled in proportion to its row count raised to
`mixture.temperature` (default 0.5). `max_examples_per_source` caps rows per
source. `loss_weight` multiplies each source's loss. Loss is averaged over the
questions on a canvas and then over canvases, so a canvas with many questions
does not outweigh a single-question canvas.
To train on a prebuilt draw sequence, set `mixture.premixed: true` and
`per_batch_stratified: false`. Each training row then needs a unique draw ID
and an origin matching its own identity; copies of one origin must be identical
apart from this metadata:
```json
"source_metadata": {
"premix": {
"draw_id": "run-000001",
"origin": {"source": "support_demo", "id": "case-0042", "group": "case-0042"}
}
}
```
### Batching
`batch_flattening: true` concatenates the logical microbatch and keeps
decisions isolated through the varlen attention path. For variable-length
decisions, set `batch_flattening: false`, `sample_packing: true` and
`eval_sample_packing: false` instead. Packing uses `sequence_len` as the
per-decision limit and `sequence_len * micro_batch_size` as the physical token
budget. It requires `mixture.per_batch_stratified: false` and does not support
sampled latent-slot counts. Reduce `micro_batch_size` and raise
`gradient_accumulation_steps` if the effective batch does not fit.
### Outputs
Each saved adapter includes `diffusion_decision_manifest.json`, which records
the base model and revision, diffusion spec, canvas width, label codebook,
latent slots, read protocol and tokenizer special IDs that an external
inference integration must reproduce. Axolotl does not ship a decision-serving
endpoint. Use a separate config for a frozen test split, and do not use test
data for checkpoint selection.
### Public procedural decision mix
`scripts/diffusion_lm/build_public_decision_mix.py` materializes a reproducible
train/dev mix from
[`tasksource/procedural-typed-decisions`](https://huggingface.co/datasets/tasksource/procedural-typed-decisions)
at revision `916e6cce365a65c37d58db70c9e369795817651e`:
```bash
python scripts/diffusion_lm/build_public_decision_mix.py ./data/public-procedural
axolotl train examples/nemotron-diffusion/decision-lora-8b-public-procedural.yaml
```
The output contains only rows from that dataset plus a `contract.json`
recording the pinned revision and licenses. The dataset card declares
Apache-2.0; the separate generator repository publishes its code under
CC-BY-4.0. Keep `contract.json` with any materialized output you share. The
model is covered by its own model-card license.
`--protected-normalized PATH` (repeatable) excludes any public row whose
canonical state matches a row in your normalized evaluation JSONL; that file is
never copied to the output. This is canonical-state decontamination only, not
cross-source group or family disjointness. The train and dev outputs are
checked against each other for state, family and group overlap.
`decision-lora-8b-public-procedural.yaml` pins the model revision and sets
`mixture.temperature: 1.0`, which schedules each usable row once. At the pinned
revisions, preprocessing keeps 30,080 train rows at the 2,048-token budget,
which is 235 effective batches of 128, hence `max_steps: 235`. This mix is a
reproducible starting point for the format, not a quality result.
## Converting a causal LM
`diffusion_lm.from_causal_lm: true` trains an ordinary causal checkpoint with
docs/support-matrix.qmd
+4 -1
View File
@@ -52,6 +52,7 @@ Most people start with two, each a single plugin to turn on: [Liger](#fused-trai
| EAFT | `use_eaft: true` (+ `eaft_alpha`, `eaft_k`) | n/a | prompt to response | SFT loss weighted by top-k entropy |
| DFT | `use_dft: true` | n/a | prompt to response | SFT loss weighted by detached token probability ([paper](https://arxiv.org/abs/2508.05629)) |
| Native diffusion LM 🟡 | `diffusion_lm:` | n/a | prompt to response | masked/uniform-noise denoising objective from the model profile; LoRA only; see [diffusion_lm](diffusion_lm.qmd) |
| Typed decisions 🟡 | plugin `diffusion_decision` + `diffusion_decision:` | n/a | typed decision JSONL | choice/score/yes-no labels on a diffusion canvas |
Most runs start with **SFT** (the default). Reach for the preference methods (DPO and its variants) once you have chosen/rejected pairs, and GRPO when you have a verifiable reward. The rest are specialized.
@@ -196,6 +197,7 @@ Generic HuggingFace support is universal. This table lists **added acceleration/
| Stepwise supervised | `stepwise_supervised` | PRM data |
| Preference | `dpo_datasets:` / `kto_datasets:` | chosen/rejected or completion+label |
| Multimodal 🟡 | `chat_template` + image/audio | resizing, role boundaries |
| Typed decisions 🟡 | `diffusion_decision.jsonl`, `.typed_decisions`, `.procedural` | `diffusion_decision` plugin only |
## Optimizers & schedulers
@@ -211,6 +213,7 @@ Generic HuggingFace support is universal. This table lists **added acceleration/
|---|---|
| `kd` | knowledge distillation (offline + online vLLM/SGLang teacher) |
| `diffusion` 🔴 | legacy `diffusion:` block; translated to `diffusion_lm.from_causal_lm: true` |
| `diffusion_decision` 🟡 | typed-decision training for native diffusion models |
| `hatchery` | remote training (Tinker/Hatchery backends) |
| `nemo_gym` | RL environments / verifiable rewards |
| `llm_compressor` | sparse fine-tuning |
@@ -289,7 +292,7 @@ Works, but flagged unstable:
- Multimodal / VLM ("limited, no full feature parity")
- N-D parallelism and Tensor Parallel
- Expert Parallel (`deep_ep` or `torch` backend)
- Native diffusion LMs (`diffusion_lm:`)
- Native diffusion LMs (`diffusion_lm:`) and typed-decision training
### 🚧 WIP / feature-branch only
Not in a tagged release yet:
examples/nemotron-diffusion/README.md
+4 -1
View File
@@ -2,11 +2,14 @@
Native diffusion recipes for Nemotron-Labs-Diffusion. See
[`docs/diffusion_lm.qmd`](../../docs/diffusion_lm.qmd) for installation,
attention, packing, objective options and adapter constraints.
attention, packing, objective options, adapter constraints and the
typed-decision dataset format.
| Config | Purpose |
|---|---|
| `lora-smoke.yaml` | Pinned 3B, two-step chat LoRA smoke with packing; not a quality recipe |
| `decision-lora-8b.yaml` | 8B typed-decision LoRA starting recipe; set your own train/dev JSONL paths |
| `decision-lora-8b-public-procedural.yaml` | 8B typed-decision LoRA on the public procedural mix from `scripts/diffusion_lm/build_public_decision_mix.py` |
```bash
axolotl preprocess examples/nemotron-diffusion/lora-smoke.yaml
examples/nemotron-diffusion/decision-lora-8b-public-procedural.yaml
+82
View File
@@ -0,0 +1,82 @@
# Data: see "Public procedural decision mix" in docs/diffusion_lm.qmd.
base_model: nvidia/Nemotron-Labs-Diffusion-8B
revision_of_model: 16c67f0560b912e93e0cabb6e0c4f5c3086d95fc
trust_remote_code: true
plugins:
- axolotl.integrations.diffusion_decision.DiffusionDecisionPlugin
adapter: lora
lora_r: 64
lora_alpha: 128
lora_dropout: 0.0
lora_target_modules: [q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj]
loraplus_lr_ratio: 8
# varlen is the only diffusion attention backend that supports batch_flattening: it runs
# bidirectional attention over the concatenated microbatch from document offsets, with no
# dense mask. flex_attention and sdpa isolate documents only via sample_packing's dense mask.
attn_implementation: varlen
bf16: true
sequence_len: 2048
batch_flattening: true
sample_packing: false
micro_batch_size: 16
gradient_accumulation_steps: 8
num_epochs: 1
max_steps: 235
learning_rate: 6.5e-6
optimizer: adamw_torch
lr_scheduler: constant
warmup_steps: 0
max_grad_norm: 1.0
gradient_checkpointing: true
gradient_checkpointing_kwargs:
use_reentrant: false
dataloader_num_workers: 0
dataset_num_proc: 4
remove_unused_columns: false
val_set_size: 0
chat_template: tokenizer_default
datasets:
- path: json
data_files: ./data/public-procedural/train.jsonl
split: train
type: diffusion_decision.jsonl
test_datasets:
- path: json
data_files: ./data/public-procedural/dev.jsonl
split: dev
type: diffusion_decision.jsonl
dataset_prepared_path: ./last_run_prepared/public-procedural
output_dir: ./outputs/public-procedural
logging_steps: 1
eval_strategy: steps
eval_steps: 50
save_strategy: steps
save_steps: 50
save_total_limit: 2
eval_on_start: true
load_best_model_at_end: true
metric_for_best_model: eval_loss
greater_is_better: false
diffusion_lm:
from_causal_lm: false
generate_samples: false
canvas_width: 128
mask_token_id: 100
eos_tail: none
unroll:
k_max: 2
grad_through_steps: false
diffusion_decision:
mixture:
temperature: 1.0
read_fraction: 1.0
latent:
mode: none
labels:
label_softmax: both
brier_weight: 0.1
codebook: vendored26
seed: 42
examples/nemotron-diffusion/decision-lora-8b.yaml
+87
View File
@@ -0,0 +1,87 @@
base_model: nvidia/Nemotron-Labs-Diffusion-8B
trust_remote_code: true
plugins:
- axolotl.integrations.diffusion_decision.DiffusionDecisionPlugin
adapter: lora
lora_r: 64
lora_alpha: 128
lora_dropout: 0.0
lora_target_modules:
- q_proj
- k_proj
- v_proj
- o_proj
- gate_proj
- up_proj
- down_proj
loraplus_lr_ratio: 8
# varlen is the only diffusion attention backend that supports batch_flattening: it runs
# bidirectional attention over the concatenated microbatch from document offsets, with no
# dense mask. flex_attention and sdpa isolate documents only via sample_packing's dense mask.
attn_implementation: varlen
bf16: true
sequence_len: 2048
batch_flattening: true
sample_packing: false
eval_sample_packing: false
micro_batch_size: 16
gradient_accumulation_steps: 8
num_epochs: 1
learning_rate: 6.5e-6
optimizer: adamw_torch
lr_scheduler: constant
warmup_steps: 0
max_grad_norm: 1.0
gradient_checkpointing: true
gradient_checkpointing_kwargs:
use_reentrant: false
dataloader_num_workers: 0
dataset_num_proc: 4
val_set_size: 0
remove_unused_columns: false
datasets:
- path: json
data_files: ./data/decision-train.jsonl
split: train
type: diffusion_decision.jsonl
test_datasets:
- path: json
data_files: ./data/decision-dev.jsonl
split: dev
type: diffusion_decision.jsonl
dataset_prepared_path: ./last_run_prepared/nemotron-decision-8b
output_dir: ./outputs/nemotron-decision-8b
logging_steps: 1
save_strategy: steps
save_steps: 50
eval_strategy: steps
eval_steps: 50
load_best_model_at_end: true
metric_for_best_model: eval_loss
greater_is_better: false
diffusion_lm:
from_causal_lm: false
generate_samples: false
canvas_width: 128
mask_token_id: 100
eos_tail: none
unroll:
k_max: 2
grad_through_steps: false
diffusion_decision:
read_fraction: 1.0
mixture:
premixed: false
per_batch_stratified: false
latent:
mode: none
labels:
label_softmax: both
brier_weight: 0.1
codebook: vendored26
seed: 42
pyproject.toml
+3 -2
View File
@@ -29,7 +29,7 @@ dependencies = [
"hf_transfer",
"sentencepiece",
"gradio==6.10.0",
"pydantic==2.12.5",
"pydantic==2.13.5",
"addict",
"fire",
"PyYAML==6.0.3",
@@ -125,7 +125,7 @@ optimizers = [
"q-galore-torch==1.0",
]
ray = [
"ray[train]>=2.52.1",
"ray[train]>=2.59.0",
]
vllm = [
"vllm>=0.17.0",
@@ -191,6 +191,7 @@ axolotl = [
"core/trainers/prm/LICENSE",
"core/trainers/prm/NOTICE",
"integrations/kernels/libs/scattermoe_lora/metadata.json",
"integrations/diffusion_decision/vendored/DJEV_LICENSE",
]
"axolotl.integrations.baseten" = ["template/*.sh", "template/*.py"]
scripts/diffusion_lm/build_public_decision_mix.py
+245
View File
@@ -0,0 +1,245 @@
"""Materialize a pinned, public-only typed-decision train/dev mix.
The generator never downloads or writes protected evaluation data. It can
optionally use locally normalized evaluation JSONL rows to reject equal states
before public output is written.
"""
from __future__ import annotations
import argparse
import hashlib
import json
from collections import Counter
from collections.abc import Iterable, Mapping, Sequence
from pathlib import Path
from typing import Any
from datasets import load_dataset
from axolotl.integrations.diffusion_decision.adapters import normalize_record
from axolotl.integrations.diffusion_decision.hygiene import (
assert_split_isolation,
decontaminate,
)
DATASET = "tasksource/procedural-typed-decisions"
REVISION = "916e6cce365a65c37d58db70c9e369795817651e"
GENERATOR_REPOSITORY_LICENSE = "CC-BY-4.0"
DATASET_CARD_LICENSE = "Apache-2.0"
CONFIGS = (
"arithmetic",
"entity_belief_tracking",
"event_state_reconstruction",
"evidence_sufficiency",
"multi_view_adjudication",
"needle_retrieval",
"partial_observation_calibration",
"policy_applicability",
"policy_under_uncertainty",
"record_aggregation",
"state_perturbation",
"table_lookup",
)
TRAIN_PER_CONFIG = 2560
DEV_PER_CONFIG = 256
SEED = "public-procedural-typed-decisions-v1"
CODEBOOK = "vendored26"
def _canonical(value: Any) -> str:
return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
def _sha256_bytes(path: Path) -> str:
return hashlib.sha256(path.read_bytes()).hexdigest()
def _read_jsonl(paths: Iterable[Path]) -> list[dict[str, Any]]:
rows: list[dict[str, Any]] = []
for path in paths:
with path.open(encoding="utf-8") as handle:
for line_number, line in enumerate(handle, start=1):
if line.strip():
try:
row = json.loads(line)
except json.JSONDecodeError as error:
raise ValueError(
f"invalid JSON in {path}:{line_number}"
) from error
if not isinstance(row, dict):
raise ValueError(
f"protected row in {path}:{line_number} is not an object"
)
rows.append(row)
return rows
def _rank(record: Mapping[str, Any], config: str) -> str:
record_id = record.get("id")
if not isinstance(record_id, str) or not record_id:
raise ValueError("public source record requires a nonempty id")
return hashlib.sha256(f"{SEED}\0{config}\0{record_id}".encode()).hexdigest()
def _normalize_row(
row: Mapping[str, Any], *, config: str, split: str
) -> dict[str, str]:
record = normalize_record("procedural", dict(row), training=True, codebook=CODEBOOK)
record["source"] = f"public_procedural.{config}"
record["family"] = f"{config}:{split}"
metadata = dict(record.get("source_metadata", {}))
metadata["provenance"] = {
"dataset": DATASET,
"revision": REVISION,
"license_evidence": {
"dataset_card_declared": DATASET_CARD_LICENSE,
"linked_generator_repository": GENERATOR_REPOSITORY_LICENSE,
},
"config": config,
"split": split,
"source_id": record["id"],
}
record["source_metadata"] = metadata
return {"normalized": _canonical(record)}
def _normalize(config: str, split: str) -> list[dict[str, Any]]:
source = load_dataset(DATASET, name=config, split=split, revision=REVISION)
normalized = source.map(
_normalize_row,
fn_kwargs={"config": config, "split": split},
num_proc=4,
remove_columns=source.column_names,
keep_in_memory=True,
desc=f"Normalizing {config} {split}",
)
return [json.loads(row["normalized"]) for row in normalized]
def _select(
records: Sequence[Mapping[str, Any]], config: str, limit: int
) -> list[Mapping[str, Any]]:
if len(records) < limit:
raise ValueError(f"{config} has {len(records)} rows; need {limit}")
return sorted(records, key=lambda record: _rank(record, config))[:limit]
def _require_unique_ids(records: list[dict[str, Any]], config: str, split: str) -> None:
ids = [record.get("id") for record in records]
if len(ids) != len(set(ids)):
raise ValueError(f"{config} {split} has duplicate public source ids")
def _write_jsonl(path: Path, rows: Iterable[Mapping[str, Any]]) -> None:
path.write_text("".join(_canonical(row) + "\n" for row in rows), encoding="utf-8")
def _assert_source_group_isolation(
splits: Mapping[str, Iterable[Mapping[str, Any]]],
) -> None:
origins: dict[tuple[str, str], str] = {}
for split, records in splits.items():
for record in records:
source = record.get("source")
group = record.get("group")
if not isinstance(source, str) or not source:
raise ValueError("decision record requires a source")
if not isinstance(group, str) or not group:
raise ValueError("decision record requires a group")
key = source, group
if key in origins and origins[key] != split:
raise ValueError(
f"source/group {key!r} overlaps splits {origins[key]!r} and {split!r}"
)
origins[key] = split
def build(output_dir: Path, protected_paths: Iterable[Path] = ()) -> dict[str, Any]:
protected_paths = list(protected_paths)
protected = _read_jsonl(protected_paths)
protected_sources = Counter(str(record.get("source", "")) for record in protected)
train, dev = [], []
for config in CONFIGS:
candidate_train = _normalize(config, "train")
candidate_dev = _normalize(config, "validation")
_require_unique_ids(candidate_train, config, "train")
_require_unique_ids(candidate_dev, config, "validation")
kept_train, _ = decontaminate(
candidate_train, protected, exclude_families=False
)
kept_dev, _ = decontaminate(candidate_dev, protected, exclude_families=False)
train.extend(_select(kept_train, config, TRAIN_PER_CONFIG))
dev.extend(_select(kept_dev, config, DEV_PER_CONFIG))
assert_split_isolation({"train": train, "dev": dev})
_assert_source_group_isolation({"train": train, "dev": dev})
output_dir.mkdir(parents=True, exist_ok=True)
train_path, dev_path = output_dir / "train.jsonl", output_dir / "dev.jsonl"
_write_jsonl(train_path, train)
_write_jsonl(dev_path, dev)
contract = {
"schema_version": 1,
"name": "public-procedural-typed-decisions-v1",
"purpose": "Public, reproducible typed-decision train/dev materialization; not a quality claim.",
"source": {
"dataset": DATASET,
"revision": REVISION,
"card": f"https://huggingface.co/datasets/{DATASET}",
"generator_repository": "https://github.com/sileod/tasksource",
"license_evidence": {
"dataset_card_declared": DATASET_CARD_LICENSE,
"linked_generator_repository": GENERATOR_REPOSITORY_LICENSE,
},
"redistribution_status": "dataset card declares Apache-2.0; retain this provenance when sharing materialized output",
"configs": list(CONFIGS),
},
"selection": {
"algorithm": "sha256(seed\\0config\\0id), ascending",
"seed": SEED,
"codebook": CODEBOOK,
"train_per_config": TRAIN_PER_CONFIG,
"dev_per_config": DEV_PER_CONFIG,
},
"protected_decontamination": {
"status": "performed" if protected else "not_requested",
"checked_when_requested": ["canonical_state"],
"source_family_group_check": "not comparable across sources: hygiene family keys and source/group keys include source; a cross-source semantic-group mapping is not available",
"protected_rows_by_source": dict(sorted(protected_sources.items())),
"protected_input_sha256": [_sha256_bytes(path) for path in protected_paths],
},
"files": {
"train": {
"path": "train.jsonl",
"lines": len(train),
"sha256": _sha256_bytes(train_path),
},
"dev": {
"path": "dev.jsonl",
"lines": len(dev),
"sha256": _sha256_bytes(dev_path),
},
},
}
(output_dir / "contract.json").write_text(
_canonical(contract) + "\n", encoding="utf-8"
)
return contract
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("output_dir", type=Path)
parser.add_argument(
"--protected-normalized",
type=Path,
action="append",
help="Optional untracked normalized evaluation JSONL for canonical-state exclusion; never copied to output.",
)
args = parser.parse_args()
build(args.output_dir, args.protected_normalized or ())
if __name__ == "__main__":
main()
src/axolotl/integrations/diffusion_decision/__init__.py
+16
View File
@@ -0,0 +1,16 @@
"""Typed decision training for native diffusion language models."""
__ci_config_keys__ = ("diffusion_decision",)
from .args import DiffusionDecisionArgs, DiffusionDecisionConfig
from .plugin import DiffusionDecisionPlugin
from .records import DecisionCanvas
from .slots import SlotInit
__all__ = [
"DecisionCanvas",
"DiffusionDecisionArgs",
"DiffusionDecisionConfig",
"DiffusionDecisionPlugin",
"SlotInit",
]
src/axolotl/integrations/diffusion_decision/_util.py
+153
View File
@@ -0,0 +1,153 @@
"""Shared config, JSON, hashing, and file helpers for decision training."""
from __future__ import annotations
import hashlib
import json
import os
import tempfile
from collections.abc import Mapping, Sequence
from pathlib import Path
from typing import Any
from axolotl.model_support import (
DiffusionSpec,
get_model_support_for_cfg,
resolve_model_support,
)
from axolotl.utils.dict import DictDefault
CONTROL_TOKEN_FIELDS = (
"pad_token_id",
"bos_token_id",
"eos_token_id",
"unk_token_id",
"mask_token_id",
)
def _value(value: Any, name: str, default: Any = None) -> Any:
if isinstance(value, Mapping):
return value.get(name, default)
return getattr(value, name, default)
def model_config_overrides(cfg: Any) -> Any:
"""``model_config`` overrides, which validated configs store under the field name."""
overrides = _value(cfg, "overrides_of_model_config")
return _value(cfg, "model_config") if overrides is None else overrides
def require_diffusion_spec(cfg: Any) -> DiffusionSpec:
if isinstance(cfg, Mapping) and not isinstance(cfg, DictDefault):
cfg = DictDefault(cfg)
profile = resolve_model_support(get_model_support_for_cfg(cfg))
spec = None if profile is None else profile.diffusion
if not isinstance(spec, DiffusionSpec):
raise ValueError("diffusion_decision requires a resolved native DiffusionSpec")
return spec
def parse_token_ids(value: Any, name: str) -> tuple[int, ...]:
if value is None:
return ()
values = (
(value,) if isinstance(value, int) and not isinstance(value, bool) else value
)
if not isinstance(values, Sequence) or isinstance(values, (str, bytes)):
raise ValueError(f"{name} must be an integer sequence")
for token_id in values:
if isinstance(token_id, bool) or not isinstance(token_id, int) or token_id < 0:
raise ValueError(f"{name} must contain nonnegative integers")
return tuple(values)
def optional_token_id(value: Any, name: str) -> int | None:
token_ids = parse_token_ids(value, name)
if len(token_ids) > 1:
raise ValueError(f"{name} must resolve to at most one token ID")
return token_ids[0] if token_ids else None
def required_token_id(value: Any, name: str) -> int:
token_id = optional_token_id(value, name)
if token_id is None:
raise ValueError(f"{name} must resolve to exactly one token ID")
return token_id
def resolve_mask_token_id(diffusion: Any, model_config: Any) -> int | None:
value = _value(diffusion, "mask_token_id")
if value is None:
value = _value(model_config, "mask_token_id")
return optional_token_id(value, "mask_token_id")
def active_control_ids(*sources: Any) -> set[int]:
return {
token_id
for source in sources
for name in CONTROL_TOKEN_FIELDS
for token_id in parse_token_ids(_value(source, name), name)
}
def canonical_json(value: Any, *, ensure_ascii: bool = False, **kwargs: Any) -> str:
return json.dumps(
value,
ensure_ascii=ensure_ascii,
sort_keys=True,
separators=(",", ":"),
**kwargs,
)
def canonical_sha256(value: Any, **kwargs: Any) -> str:
return hashlib.sha256(canonical_json(value, **kwargs).encode("utf-8")).hexdigest()
def stable_seed(seed: int, key: str) -> int:
return int.from_bytes(hashlib.sha256(f"{seed}:{key}".encode()).digest()[:8], "big")
def sha256_file(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
for chunk in iter(lambda: handle.read(1 << 20), b""):
digest.update(chunk)
return digest.hexdigest()
def read_jsonl(path: Path) -> list[dict[str, Any]]:
records: list[dict[str, Any]] = []
with path.open(encoding="utf-8") as handle:
for line_number, line in enumerate(handle, start=1):
if not line.strip():
continue
try:
record = json.loads(line)
except json.JSONDecodeError as error:
raise ValueError(
f"invalid JSONL at {path}:{line_number}: {error.msg}"
) from error
if not isinstance(record, dict):
raise ValueError(
f"JSONL record at {path}:{line_number} must be an object"
)
records.append(record)
return records
def atomic_write_text(path: Path, text: str) -> None:
descriptor, temporary = tempfile.mkstemp(
prefix=f".{path.name}.", suffix=".tmp", dir=path.parent, text=True
)
try:
with os.fdopen(descriptor, "w", encoding="utf-8") as handle:
handle.write(text)
handle.flush()
os.fsync(handle.fileno())
os.replace(temporary, path)
except BaseException:
Path(temporary).unlink(missing_ok=True)
raise
src/axolotl/integrations/diffusion_decision/adapters/__init__.py
+49
View File
@@ -0,0 +1,49 @@
"""Source adapters for typed-decision datasets."""
from __future__ import annotations
from collections.abc import Callable, Mapping
from typing import Any
from .jsonl import normalize_jsonl
from .procedural import normalize as normalize_procedural
from .typed_decisions import normalize as normalize_typed_decisions
_Adapter = Callable[[Mapping[str, Any], str | None], Mapping[str, Any]]
_ADAPTERS: dict[str, _Adapter] = {
"jsonl": lambda row, _split: row,
"procedural": lambda row, _split: normalize_procedural(row),
"typed_decisions": lambda row, split: normalize_typed_decisions(
row, source_split=split
),
}
def normalize_record(
adapter: str,
row: Mapping[str, Any],
*,
training: bool,
codebook: str = "vendored26",
source_split: str | None = None,
) -> dict[str, Any]:
name = adapter.removeprefix("diffusion_decision.")
if name not in _ADAPTERS:
raise ValueError(f"unknown decision adapter: {adapter}")
if training and name == "typed_decisions" and source_split != "train":
raise ValueError("typed_decisions training requires the declared train split")
result = normalize_jsonl(_ADAPTERS[name](row, source_split), codebook=codebook)
if training and name == "jsonl" and result.get("source") == "typed_decisions":
raise ValueError(
"typed_decisions is reserved for evaluation when supplied through jsonl"
)
return result
__all__ = [
"normalize_record",
"normalize_procedural",
"normalize_typed_decisions",
"normalize_jsonl",
]
src/axolotl/integrations/diffusion_decision/adapters/common.py
+91
View File
@@ -0,0 +1,91 @@
"""Shared validation for normalized typed-decision source records."""
from __future__ import annotations
import json
import math
from collections.abc import Sequence
from typing import Any, Literal
LabelKind = Literal["hard", "dist", "set"]
def decode_json(value: Any) -> Any:
"""Decode structured source fields while retaining ordinary text states."""
if not isinstance(value, str):
return value
try:
return json.loads(value)
except json.JSONDecodeError:
return value
def probabilities(values: Sequence[float], count: int) -> list[float]:
if len(values) != count or count < 2:
raise ValueError("target must contain one probability per alternative")
result = [float(value) for value in values]
if any(not math.isfinite(value) or value < 0 for value in result):
raise ValueError("target probabilities must be finite and nonnegative")
total = sum(result)
if not math.isclose(total, 1.0, rel_tol=0, abs_tol=1e-4):
raise ValueError("target probabilities must sum to one")
return [value / total for value in result]
def label_target(
values: Sequence[float],
kind: LabelKind,
*,
candidate_ids: Sequence[str] | None = None,
) -> dict[str, Any]:
support = [index for index, value in enumerate(values) if value > 0]
if kind == "dist":
if candidate_ids is None:
return {"kind": "dist", "probs": list(values)}
if len(candidate_ids) != len(values):
raise ValueError("soft candidate IDs must align with probabilities")
return {
"kind": "dist",
"candidate_indices": support,
"candidate_ids": [str(candidate_ids[index]) for index in support],
"probs": [values[index] for index in support],
"other_probability": sum(
value for index, value in enumerate(values) if index not in support
),
}
if kind == "set":
if not support:
raise ValueError("set target must contain an allowed alternative")
return {"kind": "set", "allowed_set": support}
if kind != "hard":
raise ValueError(f"unknown target kind: {kind}")
if len(support) != 1:
raise ValueError("hard target must be one-hot; refusing to discard soft mass")
return {"kind": "hard", "gold_idx": support[0]}
def finite_question(
kind: str, instructions: str, options: Sequence[str], target: Sequence[float]
) -> tuple[dict[str, Any], list[float]]:
"""Normalize a question and remap its target to the serving option order."""
names = [str(option) for option in options]
if len(set(names)) != len(names):
raise ValueError("question alternatives must have distinct names")
values = probabilities(target, len(names))
question: dict[str, Any] = {"type": kind, "instructions": instructions}
if kind in {"noul", "bool", "boolean"}:
aliases = {"true": "yes", "false": "no", "yes": "yes", "no": "no"}
canonical = [aliases.get(name.lower()) for name in names]
if sorted(name for name in canonical if name is not None) != ["no", "yes"]:
raise ValueError(
"boolean alternatives must be exactly yes/no or true/false"
)
question["type"] = "noul"
values = [values[canonical.index(name)] for name in ("yes", "no")]
elif kind == "choice":
question["options"] = names
elif kind == "score":
question["levels"] = names
else:
raise ValueError(f"unsupported finite question type: {kind}")
return question, values
src/axolotl/integrations/diffusion_decision/adapters/jsonl.py
+111
View File
@@ -0,0 +1,111 @@
"""Validate user-supplied records in the normalized decision format."""
from __future__ import annotations
import math
from collections.abc import Mapping
from copy import deepcopy
from typing import Any
from ..template import parse_decision_schema
from .common import probabilities
def normalize_jsonl(
row: Mapping[str, Any], *, codebook: str = "vendored26"
) -> dict[str, Any]:
result = deepcopy(dict(row))
for field in ("source", "group"):
if not isinstance(result.get(field), str) or not result[field]:
raise ValueError(f"normalized record requires {field}")
if result.get("state") is None:
raise ValueError("normalized record requires state")
questions, labels = result.get("questions"), result.get("labels")
if not isinstance(questions, dict) or not isinstance(labels, dict):
raise ValueError("questions and labels must be objects")
if set(questions) != set(labels):
raise ValueError("each question must have exactly one target")
parsed = parse_decision_schema(
{"questions": [{**q, "id": key} for key, q in questions.items()]},
codebook=codebook,
)
for question in parsed["questions"]:
key = question["id"]
if question["type"] not in {"noul", "choice", "score"}:
raise ValueError("decision training requires finite label alternatives")
count = len(question["labels"])
target = labels[key]
if not isinstance(target, dict):
raise ValueError(f"question {key!r} target must be an object")
kind = target.get("kind")
if kind == "dist":
indices = target.get("candidate_indices")
if indices is None:
target["probs"] = probabilities(target["probs"], count)
else:
if (
not isinstance(indices, list)
or not indices
or len(indices) != len(target.get("probs", ()))
or len(set(indices)) != len(indices)
or any(
isinstance(index, bool)
or not isinstance(index, int)
or not 0 <= index < count
for index in indices
)
):
raise ValueError(
f"question {key!r} sparse candidates must identify alternatives"
)
values = [float(value) for value in target["probs"]]
other = float(target.get("other_probability", 0.0))
if (
any(not math.isfinite(value) or value < 0 for value in values)
or not math.isfinite(other)
or other < 0
or not math.isclose(sum(values) + other, 1.0, abs_tol=1e-4)
):
raise ValueError(
f"question {key!r} sparse probabilities plus "
"other_probability must sum to one"
)
target["probs"] = values
target["other_probability"] = other
candidate_ids = target.get("candidate_ids")
if candidate_ids is not None and (
not isinstance(candidate_ids, list)
or len(candidate_ids) != len(indices)
or any(
not isinstance(value, str) or not value
for value in candidate_ids
)
):
raise ValueError(
f"question {key!r} sparse candidate_ids must align "
"with candidates"
)
elif kind in {"hard", "set"}:
indices = (
[target.get("gold_idx")]
if kind == "hard"
else target.get("allowed_set")
)
if (
not isinstance(indices, list)
or not indices
or any(
isinstance(index, bool)
or not isinstance(index, int)
or not 0 <= index < count
for index in indices
)
):
raise ValueError(
f"question {key!r} target indices must identify alternatives"
)
if len(set(indices)) != len(indices):
raise ValueError(f"question {key!r} allowed set must be unique")
else:
raise ValueError(f"unknown label kind {kind!r}")
return result
src/axolotl/integrations/diffusion_decision/adapters/procedural.py
+82
View File
@@ -0,0 +1,82 @@
from __future__ import annotations
from .common import decode_json, finite_question, label_target
def normalize(row):
questions, answers = decode_json(row["questions"]), decode_json(row["answers"])
out = {}
labels = {}
for key, q in questions.items():
if key not in answers:
raise ValueError(f"missing target for {key}")
a = answers[key]
kind = q["type"]
criteria = q.get("criteria", {})
if kind == "choice":
options = list(criteria)
descriptions = [criteria[x] for x in options]
elif kind == "score":
options = [str(x) for x in criteria]
descriptions = options
elif kind == "noul":
options = ["yes", "no"]
descriptions = (
[criteria.get("true"), criteria.get("false")]
if isinstance(criteria, dict)
else [None, None]
)
else:
raise ValueError(f"unsupported procedural type {kind}")
if kind == "noul":
if "noul" not in a:
raise ValueError(f"missing noul target for {key}")
values = [float(a["noul"]), 1 - float(a["noul"])]
else:
probs = a.get("probabilities")
if not isinstance(probs, dict):
raise ValueError(f"incomplete probabilities for {key}")
identities = options
if kind == "score" and a.get("legend") is not None:
legend = a["legend"]
if not isinstance(legend, dict) or set(legend) != set(probs):
raise ValueError(
f"score legend must identify every probability for {key}"
)
inverse = {str(level): identity for identity, level in legend.items()}
if len(inverse) != len(legend) or set(inverse) != set(options):
raise ValueError(
f"score legend must identify each level exactly once for {key}"
)
identities = [inverse[level] for level in options]
elif kind == "score" and set(probs) != set(options):
identities = [str(index) for index in range(len(options))]
if set(probs) != set(identities):
raise ValueError(f"incomplete probabilities for {key}")
values = [float(probs[identity]) for identity in identities]
question, values = finite_question(
kind, q.get("instructions", ""), options, values
)
if kind == "choice":
question["options"] = [
{"name": name, "description": desc}
for name, desc in zip(options, descriptions, strict=True)
]
elif kind == "noul":
question["criteria"] = {"true": descriptions[0], "false": descriptions[1]}
labels[key] = label_target(values, "hard" if max(values) == 1 else "dist")
out[key] = question
record_id = str(row["id"])
task, separator, official_split = record_id.partition(":")
metadata = {}
if separator and official_split:
metadata = {"task": task, "official_split": official_split.split(":", 1)[0]}
return {
"source": "procedural",
"id": record_id,
"group": record_id,
"source_metadata": metadata,
"state": decode_json(row["state"]),
"questions": out,
"labels": labels,
}
src/axolotl/integrations/diffusion_decision/adapters/typed_decisions.py
+103
View File
@@ -0,0 +1,103 @@
"""Normalize the multi-question typed-decisions evaluation corpus."""
from __future__ import annotations
import math
from collections.abc import Mapping
from typing import Any
from .common import decode_json, finite_question, label_target
def normalize(
row: Mapping[str, Any], *, source_split: str | None = None
) -> dict[str, Any]:
row_split = row.get("split")
if row_split is not None and (
not isinstance(row_split, str) or row_split.lower() != source_split
):
raise ValueError("typed decisions row split disagrees with declared split")
questions, gold = decode_json(row["questions"]), decode_json(row["gold"])
if not isinstance(questions, dict) or not questions or not isinstance(gold, dict):
raise ValueError("typed decisions require questions and gold objects")
if set(questions) != set(gold):
raise ValueError("each typed question must have exactly one target")
normalized, labels = {}, {}
for key, question in questions.items():
answer = gold[key]
kind, criteria = question["type"], question.get("criteria")
if kind == "choice":
if not isinstance(criteria, dict):
raise ValueError("choice criteria must map names to descriptions")
identities = list(criteria)
names = identities
elif kind == "score":
if not isinstance(criteria, list):
raise ValueError("score criteria must be an ordered list")
identities = [str(index) for index in range(len(criteria))]
names = [str(value) for value in criteria]
elif kind == "noul":
identities, names = ["false", "true"], ["no", "yes"]
else:
raise ValueError(f"unsupported typed question type: {kind}")
probs = answer.get("probabilities")
if not isinstance(probs, dict) or set(probs) != set(identities):
raise ValueError(
f"question {key!r} needs probabilities for each alternative"
)
result, values = finite_question(
kind,
str(question.get("instructions", "")),
names,
[probs[identity] for identity in identities],
)
if kind == "noul":
identities = ["yes", "no"]
if kind == "choice":
assert isinstance(criteria, dict)
result["options"] = [
{"name": name, "description": criteria[name]} for name in names
]
elif kind == "noul" and criteria:
result["criteria"] = criteria
normalized[key] = result
smoothing = answer.get("smoothing")
if smoothing is None:
labels[key] = label_target(values, "dist", candidate_ids=identities)
else:
if (
isinstance(smoothing, bool)
or not isinstance(smoothing, (float, int))
or not math.isfinite(smoothing)
or not 0.0 <= smoothing < 1.0
):
raise ValueError(
f"question {key!r} hard smoothing must be finite and in [0, 1)"
)
target = label_target(values, "hard")
target["smoothing"] = float(smoothing)
labels[key] = target
state = decode_json(row["state"])
if state is None:
raise ValueError("typed record requires a state")
workflow = row["workflow"]
if not isinstance(workflow, str) or not workflow:
raise ValueError("typed record requires workflow metadata")
return {
"source": "typed_decisions",
"family": workflow,
"id": str(row["id"]),
"group": str(row["id"]),
"state": state,
"questions": normalized,
"labels": labels,
"source_metadata": {
"provenance": {
"dataset": "LocalLLaMA/typed-decisions",
"workflow": workflow,
"split": source_split,
"row_split": row_split,
"source_id": str(row["id"]),
}
},
}
src/axolotl/integrations/diffusion_decision/args.py
+289
View File
@@ -0,0 +1,289 @@
"""Configuration models for typed diffusion-decision training."""
from __future__ import annotations
import math
from typing import Literal
from pydantic import BaseModel, ConfigDict, Field, model_validator
from .vendored.djev_template import MAX_QUESTIONS
class DecisionMixtureConfig(BaseModel):
"""Deterministic source-mixture controls."""
model_config = ConfigDict(extra="forbid")
weights: dict[str, float] | None = Field(
default=None,
json_schema_extra={
"description": "Per-source sampling weights keyed by record source; exclusive with temperature."
},
)
temperature: float | None = Field(
default=None,
ge=0.0,
json_schema_extra={
"description": "Sample each source in proportion to its row count raised to this power; unset uses 0.5."
},
)
per_batch_stratified: bool = Field(
default=True,
json_schema_extra={
"description": "Draw each microbatch with source-stratified sampling; must be false for packing or premixed data."
},
)
max_examples_per_source: int | dict[str, int] | None = Field(
default=None,
json_schema_extra={
"description": "Cap on training rows kept per source, as one integer or a per-source mapping."
},
)
loss_weight: dict[str, float] = Field(
default_factory=dict,
json_schema_extra={"description": "Per-source multiplier on each canvas loss."},
)
premixed: bool = Field(
default=False,
json_schema_extra={
"description": "Treat training rows as a prebuilt draw sequence; each row needs source_metadata.premix."
},
)
@model_validator(mode="after")
def validate_weights(self) -> "DecisionMixtureConfig":
if self.weights is not None and self.temperature is not None:
raise ValueError(
"mixture specifies explicit weights or temperature, not both"
)
if self.temperature is not None and not math.isfinite(self.temperature):
raise ValueError("mixture temperature must be finite")
for field_name, values in (
("weights", self.weights or {}),
("loss_weight", self.loss_weight),
):
for source, value in values.items():
if not source:
raise ValueError(
f"mixture {field_name} source names must be nonempty"
)
if not math.isfinite(value) or value <= 0:
raise ValueError(
f"mixture {field_name} values must be finite and positive"
)
if isinstance(self.max_examples_per_source, int):
if self.max_examples_per_source < 1:
raise ValueError("mixture max_examples_per_source must be positive")
elif isinstance(self.max_examples_per_source, dict):
for source, value in self.max_examples_per_source.items():
if (
not source
or isinstance(value, bool)
or not isinstance(value, int)
or value < 1
):
raise ValueError(
"mixture max_examples_per_source requires nonempty names and positive caps"
)
if self.premixed:
incompatible = []
if self.weights is not None:
incompatible.append("weights")
if self.temperature is not None:
incompatible.append("temperature")
if self.max_examples_per_source is not None:
incompatible.append("max_examples_per_source")
if self.per_batch_stratified:
incompatible.append("per_batch_stratified")
if incompatible:
raise ValueError(
"mixture premixed=true cannot be combined with "
+ ", ".join(incompatible)
)
return self
class DecisionLabelsConfig(BaseModel):
"""Typed-label objective controls."""
model_config = ConfigDict(extra="forbid")
label_softmax: Literal["restricted", "full", "both"] = Field(
default="both",
json_schema_extra={
"description": "Label CE over the question candidates, the full vocabulary, or both."
},
)
full_ce_weighting: Literal["ce", "dft"] = Field(
default="ce",
json_schema_extra={
"description": "Full-vocabulary CE for hard labels: plain CE or DFT probability weighting."
},
)
brier_weight: float = Field(
default=0.1,
ge=0.0,
json_schema_extra={
"description": "Coefficient of the candidate-restricted Brier auxiliary loss."
},
)
hard_label_smoothing: float = Field(
default=0.0,
ge=0.0,
lt=1.0,
json_schema_extra={
"description": "Fraction of a one-hot CE target spread uniformly over the valid answer tokens."
},
)
codebook: Literal["vendored26", "expanded52"] = Field(
default="vendored26",
json_schema_extra={
"description": "Answer-label codebook: djev letters A-Z (26 options) or A-Z plus a-z (52)."
},
)
@model_validator(mode="after")
def validate_brier_weight(self) -> "DecisionLabelsConfig":
if not math.isfinite(self.brier_weight):
raise ValueError("labels brier_weight must be finite")
if self.hard_label_smoothing and self.full_ce_weighting == "dft":
raise ValueError("hard_label_smoothing requires full_ce_weighting=ce")
return self
class DecisionLatentConfig(BaseModel):
"""Latent-slot initialization controls retained in the run manifest."""
model_config = ConfigDict(extra="forbid")
mode: Literal["none", "pad", "pinned", "learned", "free", "mask", "prompt"] = Field(
default="none",
json_schema_extra={
"description": "Latent-slot initialization; none adds no slots to the canvas."
},
)
num_slots: int = Field(
default=0,
ge=0,
json_schema_extra={"description": "Number of latent slots per canvas."},
)
sample_num_slots: bool = Field(
default=False,
json_schema_extra={
"description": "Sample a slot count up to num_slots per canvas instead of a fixed count."
},
)
token_ids: list[int] = Field(
default_factory=list,
json_schema_extra={
"description": "Slot token IDs for the pinned, learned, and prompt modes."
},
)
free_update_policy: Literal["argmax"] | None = Field(
default=None,
json_schema_extra={
"description": "Slot update rule between reads; required for mode free."
},
)
@model_validator(mode="after")
def validate_slots(self) -> "DecisionLatentConfig":
if any(token_id < 0 for token_id in self.token_ids):
raise ValueError("latent token_ids must be nonnegative")
if self.mode == "none" and (self.num_slots or self.token_ids):
raise ValueError("latent mode=none requires num_slots=0 and no token_ids")
if self.sample_num_slots and self.mode == "none":
raise ValueError("sample_num_slots requires a non-none latent mode")
if self.sample_num_slots and self.num_slots < 1:
raise ValueError("sample_num_slots requires num_slots to be positive")
if self.mode == "free" and self.free_update_policy != "argmax":
raise ValueError("latent mode=free requires free_update_policy=argmax")
if self.mode != "free" and self.free_update_policy is not None:
raise ValueError("free_update_policy is valid only for latent mode=free")
return self
class DecisionCarryConfig(BaseModel):
"""Cross-read latent carry policy."""
model_config = ConfigDict(extra="forbid")
enabled: bool = Field(
default=False,
json_schema_extra={
"description": "Reserved; carrying latent state across reads is not supported in training."
},
)
@model_validator(mode="after")
def reject_enabled(self) -> "DecisionCarryConfig":
if self.enabled:
raise ValueError("decision carry is not wired for training")
return self
class DiffusionDecisionConfig(BaseModel):
"""Typed-decision data, latent, and readout settings."""
model_config = ConfigDict(extra="forbid")
layout: Literal["thought_block", "prompt_slots"] = Field(
default="thought_block",
json_schema_extra={
"description": "Decision layout recorded in the adapter manifest."
},
)
max_questions_per_canvas: int = Field(
default=20,
ge=1,
le=MAX_QUESTIONS,
json_schema_extra={
"description": "Maximum questions grouped into one canvas; capped by the template."
},
)
read_fraction: float = Field(
default=1.0,
json_schema_extra={
"description": "Fraction of canvases trained at the fully masked read time; the rest use sampled diffusion times."
},
)
mixture: DecisionMixtureConfig = Field(
default_factory=DecisionMixtureConfig,
json_schema_extra={"description": "Source sampling and loss weighting."},
)
labels: DecisionLabelsConfig = Field(
default_factory=DecisionLabelsConfig,
json_schema_extra={"description": "Label objective and answer codebook."},
)
latent: DecisionLatentConfig = Field(
default_factory=DecisionLatentConfig,
json_schema_extra={
"description": "Optional latent slots placed in the canvas."
},
)
carry: DecisionCarryConfig = Field(
default_factory=DecisionCarryConfig,
json_schema_extra={
"description": "Cross-read latent carry; must stay disabled."
},
)
@model_validator(mode="after")
def validate_read_fraction(self) -> "DiffusionDecisionConfig":
if not math.isfinite(self.read_fraction):
raise ValueError("read_fraction must be finite")
if not 0.0 <= self.read_fraction <= 1.0:
raise ValueError("read_fraction must be in [0, 1]")
return self
class DiffusionDecisionArgs(BaseModel):
"""Plugin entry that exposes the optional ``diffusion_decision`` block."""
diffusion_decision: DiffusionDecisionConfig | None = Field(
default=None,
json_schema_extra={
"description": "Typed-decision training settings for native diffusion models."
},
)
src/axolotl/integrations/diffusion_decision/budgets.py
+37
View File
@@ -0,0 +1,37 @@
"""Pre-mixture token-budget checks for fixed-width decision canvases."""
from __future__ import annotations
from dataclasses import dataclass
from axolotl.model_support import DiffusionLayout
@dataclass(frozen=True)
class DecisionBudget:
logical_tokens: int
physical_tokens: int
def decision_budget(
*, prompt_tokens: int, canvas_tokens: int, layout: DiffusionLayout
) -> DecisionBudget:
if prompt_tokens < 0 or canvas_tokens <= 0:
raise ValueError("decision prompt and canvas lengths must be nonnegative")
if layout not in {DiffusionLayout.FULL_SEQUENCE, DiffusionLayout.ENCODER_CANVAS}:
raise ValueError(f"unsupported decision layout: {layout}")
total = prompt_tokens + canvas_tokens
return DecisionBudget(logical_tokens=total, physical_tokens=total)
def exceeds_budget(
budget: DecisionBudget,
*,
logical_limit: int | None,
physical_limit: int | None,
) -> str | None:
if logical_limit is not None and budget.logical_tokens > logical_limit:
return "logical"
if physical_limit is not None and budget.physical_tokens > physical_limit:
return "physical"
return None
src/axolotl/integrations/diffusion_decision/collator.py
+58
View File
@@ -0,0 +1,58 @@
from __future__ import annotations
import torch
from axolotl.core.trainers.diffusion_lm.batch import DiffusionBatch
from .records import DecisionCanvas
class DecisionCanvasCollator:
def __init__(self, pad_token_id: int = 0):
self.pad_token_id = pad_token_id
def __call__(self, canvases: list[DecisionCanvas]) -> DiffusionBatch:
if not canvases:
raise ValueError("empty decision batch")
pw = max(len(c.prompt_ids) for c in canvases)
cw = len(canvases[0].canvas_ids)
if any(len(c.canvas_ids) != cw or len(c.semantic_mask) != cw for c in canvases):
raise ValueError("canvases must share width")
b = len(canvases)
prompt = torch.full((b, pw), self.pad_token_id, dtype=torch.long)
canvas = torch.tensor([c.canvas_ids for c in canvases], dtype=torch.long)
ev = torch.zeros((b, pw), dtype=torch.bool)
ar = ev.clone()
sem = torch.tensor([c.semantic_mask for c in canvases], dtype=torch.bool)
loss = torch.zeros((b, cw), dtype=torch.bool)
for i, c in enumerate(canvases):
n = len(c.prompt_ids)
prompt[i, :n] = torch.tensor(c.prompt_ids)
ev[i, :n] = True
ar[i, 1:n] = True
if c.prompt_slot_mask:
if len(c.prompt_slot_mask) != n:
raise ValueError("prompt_slot_mask must align with prompt_ids")
ar[i, :n] &= ~torch.tensor(c.prompt_slot_mask, dtype=torch.bool)
loss[i, torch.tensor(c.label_positions, dtype=torch.long)] = True
pinned = torch.tensor([c.pinned_mask for c in canvases], dtype=torch.bool)
return DiffusionBatch(
encoder_input_ids=prompt,
encoder_validity=ev,
encoder_ar_valid_mask=ar,
encoder_document_ids=torch.arange(b)[:, None].expand(-1, pw),
encoder_position_ids=torch.arange(pw)[None].expand(b, -1),
canvas_clean_ids=canvas,
canvas_semantic_validity=sem,
canvas_loss_mask=loss,
canvas_corruptible_mask=loss,
canvas_input_pinned_mask=pinned,
canvas_sc_eligible_mask=loss & ~pinned,
canvas_read_only_mask=~loss,
canvas_update_mask=loss & ~pinned,
logical_ids=torch.arange(b),
encoder_lengths=ev.sum(1),
canvas_lengths=sem.sum(1),
decoder_prefix_lengths=ev.sum(1),
selected_block_ids=torch.zeros(b, dtype=torch.long),
)
src/axolotl/integrations/diffusion_decision/datasets.py
+812
View File
@@ -0,0 +1,812 @@
"""Dataset preparation for the model-agnostic diffusion-decision plugin."""
from __future__ import annotations
import hashlib
import logging
from collections import defaultdict
from collections.abc import Mapping, Sequence
from concurrent.futures import ThreadPoolExecutor
from functools import partial
from glob import glob
from pathlib import Path
from typing import Any
from urllib.parse import urlparse
from datasets import load_dataset
from torch.utils.data import Dataset
from axolotl.common.datasets import TrainDatasetMeta
from axolotl.core.trainers.diffusion_lm.sampling import resolve_native_packing_budget
from axolotl.loaders import load_tokenizer
from axolotl.model_support import (
DiffusionLayout,
DiffusionNoise,
)
from axolotl.utils.dict import DictDefault
from ._util import (
_value,
active_control_ids,
canonical_json,
model_config_overrides,
parse_token_ids,
read_jsonl,
require_diffusion_spec,
required_token_id,
resolve_mask_token_id,
)
from .adapters import normalize_record
from .budgets import decision_budget, exceeds_budget
from .grouping import group_records
from .hygiene import assert_split_isolation, decontaminate, family_key
from .loss import decision_example_from_canvas
from .mixture import (
DeterministicMixtureSampler,
MixtureSample,
prepare_source_pools,
)
from .permutation import permute_record
from .prepared_cache import (
identity as prepared_cache_identity,
load as load_prepared_cache,
store as store_prepared_cache,
)
from .preprocessing import build_decision_canvas
from .slot_sampling import (
DecisionDraw,
project_max_canvas,
project_slot_plan,
sample_slot_count,
)
from .slots import SlotInit, SlotPlan
from .template import SchemaError
LOG = logging.getLogger(__name__)
def _rows(entry: Any) -> list[dict[str, Any]]:
path = _value(entry, "path")
if not isinstance(path, str) or not path:
raise ValueError("decision dataset requires a path")
split = str(_value(entry, "split", "train"))
kwargs = {
key: _value(entry, key)
for key in ("name", "revision")
if _value(entry, key) is not None
}
data_files = _value(entry, "data_files")
local_jsonl = _local_jsonl_paths(path, data_files, split)
if local_jsonl is not None:
return [record for path in local_jsonl for record in read_jsonl(path)]
if (
data_files is not None
and path in {"json", "parquet", "csv", "text"}
and not isinstance(data_files, Mapping)
):
kwargs["data_files"] = {split: data_files}
elif data_files is not None:
kwargs["data_files"] = data_files
dataset = load_dataset(path, split=split, **kwargs)
return [dict(row) for row in dataset]
def _local_jsonl_paths(path: str, data_files: Any, split: str) -> list[Path] | None:
if path != "json" or data_files is None:
return None
selected = data_files.get(split) if isinstance(data_files, Mapping) else data_files
if isinstance(selected, str):
names = [selected]
elif isinstance(selected, Sequence) and not isinstance(selected, (str, bytes)):
names = list(selected)
else:
return None
if not names or not all(isinstance(name, str) for name in names):
raise ValueError("local JSONL data_files must be a nonempty path or path list")
expanded: list[Path] = []
for name in names:
if urlparse(name).scheme:
return None
matches = sorted(Path(match) for match in glob(name))
if matches:
if any(match.suffix != ".jsonl" for match in matches):
return None
expanded.extend(matches)
elif Path(name).suffix == ".jsonl":
raise ValueError(f"local JSONL path matched no files: {name}")
else:
return None
if not expanded or not all(item.is_file() for item in expanded):
raise ValueError("local JSONL data_files must resolve to regular files")
return expanded
def _prepared_cache_source_paths(entries: Sequence[Any]) -> list[Path] | None:
paths: list[Path] = []
for entry in entries:
resolved = _local_jsonl_paths(
_value(entry, "path"),
_value(entry, "data_files"),
str(_value(entry, "split", "train")),
)
if resolved is None:
return None
paths.extend(resolved)
return paths
def _prepared_cache_eligible(cfg: Any) -> bool:
latent = _value(_value(cfg, "diffusion_decision"), "latent")
return latent is None or _value(latent, "mode", "none") in {"none", "pad", "mask"}
def _premix_identity(record: Mapping[str, Any]) -> tuple[str, str, str]:
return (
str(record.get("source", "")),
str(record.get("id", "")),
str(record.get("group", "")),
)
def _validate_premixed_records(records: Sequence[Mapping[str, Any]]) -> None:
draw_origins: dict[str, str] = {}
content_by_origin: dict[tuple[str, str, str], str] = {}
for record in records:
metadata = record.get("source_metadata")
premix = metadata.get("premix") if isinstance(metadata, Mapping) else None
if not isinstance(premix, Mapping):
raise ValueError("premixed decision rows require source_metadata.premix")
draw_id = premix.get("draw_id")
origin = premix.get("origin")
if not isinstance(draw_id, str) or not draw_id:
raise ValueError("premixed decision rows require a nonempty premix draw_id")
if not isinstance(origin, Mapping) or not origin:
raise ValueError("premixed decision rows require a premix origin mapping")
identity = _premix_identity(record)
if (
tuple(str(origin.get(field, "")) for field in ("source", "id", "group"))
!= identity
):
raise ValueError(
"premixed decision origin must match source, id, and group"
)
origin_value = canonical_json(origin, ensure_ascii=True)
if draw_id in draw_origins:
raise ValueError(f"premixed decision draw_id is duplicated: {draw_id}")
draw_origins[draw_id] = origin_value
content = dict(record)
content_metadata = dict(metadata) if isinstance(metadata, Mapping) else {}
content_metadata.pop("premix", None)
content["source_metadata"] = content_metadata
canonical_content = canonical_json(content, ensure_ascii=True, default=str)
previous = content_by_origin.setdefault(identity, canonical_content)
if previous != canonical_content:
raise ValueError(
"premixed copies of one decision record must have identical content"
)
def _is_eval(entry: Any) -> bool:
return str(_value(entry, "split", "train")).lower() in {
"eval",
"test",
"validation",
"dev",
"calibration",
"ood",
}
def _family_dev_split(records: Sequence[dict[str, Any]], ratio: float = 0.2):
groups: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list)
for record in records:
groups[family_key(record)].append(record)
selected: set[tuple[str, str]] = set()
target = round(len(records) * ratio)
total = 0
for key in sorted(
groups,
key=lambda item: hashlib.sha256(f"decision-dev-v1:{item}".encode()).hexdigest(),
):
if total >= target:
break
size = len(groups[key])
if total + size >= len(records):
continue
selected.add(key)
total += size
if records:
LOG.info(
"Automatic decision dev split holds out %s of %s records",
total,
len(records),
)
return (
[record for record in records if family_key(record) not in selected],
[record for record in records if family_key(record) in selected],
)
def _thought_delimiters(model_config: Any) -> tuple[tuple[int, ...], tuple[int, ...]]:
return (
parse_token_ids(_value(model_config, "thought_open_ids"), "thought_open_ids"),
parse_token_ids(_value(model_config, "thought_close_ids"), "thought_close_ids"),
)
def _resolve_slot_plan(
tokenizer: Any, cfg: Any, spec: Any, vocab_size: int
) -> tuple[SlotPlan | None, tuple[int, ...], tuple[int, ...]]:
decision = _value(cfg, "diffusion_decision")
latent = _value(decision, "latent")
if latent is None:
return None, (), ()
mode = _value(latent, "mode", "none")
if mode == "none":
return None, (), ()
model_config = model_config_overrides(cfg)
thought_open_ids, thought_close_ids = _thought_delimiters(model_config)
token_ids = parse_token_ids(
_value(latent, "token_ids", ()) or (), "latent.token_ids"
)
mask_token_id = resolve_mask_token_id(_value(cfg, "diffusion_lm"), model_config)
if mode in {"pinned", "learned", "prompt"}:
control_ids = active_control_ids(tokenizer, model_config)
control_ids.update(thought_open_ids, thought_close_ids)
if mask_token_id is not None:
control_ids.add(mask_token_id)
collisions = sorted(set(token_ids) & control_ids)
if collisions:
raise ValueError(
"fixed decision slot token_ids cannot use active control IDs: "
f"{collisions}"
)
plan = SlotInit(
mode=mode,
token_ids=token_ids,
num_slots=_value(latent, "num_slots", 0),
vocab_size=vocab_size,
pad_id=required_token_id(
_value(tokenizer, "pad_token_id"), "tokenizer.pad_token_id"
),
spec=spec,
mask_token_id=mask_token_id,
).build(seed=int(_value(cfg, "seed", 0)))
if (
plan.placement == "thought"
and _value(spec, "layout") is DiffusionLayout.ENCODER_CANVAS
and (not thought_open_ids or not thought_close_ids)
):
raise ValueError(
"encoder-canvas thought slots require `model_config: {thought_open_ids: "
"[...], thought_close_ids: [...]}` in the training config"
)
return plan, thought_open_ids, thought_close_ids
def _canvas_row(
tokenizer,
record: dict[str, Any],
cfg: Any,
source_weight: float,
*,
spec,
vocab_size: int | None = None,
slot_plan: SlotPlan | None = None,
thought_open_ids: Sequence[int] = (),
thought_close_ids: Sequence[int] = (),
):
diffusion = _value(cfg, "diffusion_lm")
decision = _value(cfg, "diffusion_decision")
model_config = model_config_overrides(cfg)
latent = _value(decision, "latent")
configured_width = _value(diffusion, "canvas_width")
width = 128 if configured_width is None else int(configured_width)
if spec.max_canvas is not None:
width = min(width, spec.max_canvas)
mask_token_id = resolve_mask_token_id(diffusion, model_config)
scaffold_ids = (
tuple(int(token) for token in (_value(latent, "token_ids", ()) or ()))
if slot_plan is None
else ()
)
record = permute_record(
record,
seed=int(_value(cfg, "seed", 0)),
codebook=_value(_value(decision, "labels"), "codebook", "vendored26"),
)
if vocab_size is None:
vocab_size = _value(model_config, "vocab_size")
if vocab_size is None:
vocab_size = len(tokenizer)
turn_close_id = required_token_id(
_value(model_config, "turn_close_token_id", _value(tokenizer, "eos_token_id")),
"turn_close_token_id",
)
canvas = build_decision_canvas(
tokenizer,
record,
scaffold_ids=scaffold_ids,
turn_close_id=turn_close_id,
pad_id=required_token_id(
_value(tokenizer, "pad_token_id"), "tokenizer.pad_token_id"
),
vocab_size=int(vocab_size),
width=width,
seed=int(_value(cfg, "seed", 0)),
steps=int(_value(_value(diffusion, "unroll"), "k_max", 1)),
noise_kind="absorbing" if spec.noise is DiffusionNoise.ABSORBING else "uniform",
mask_token_id=mask_token_id,
slot_plan=slot_plan,
thought_open_ids=thought_open_ids,
thought_close_ids=thought_close_ids,
codebook=_value(_value(decision, "labels"), "codebook", "vendored26"),
prevalidated_record=True,
)
row = {
"canvas": canvas,
"decision_example": decision_example_from_canvas(
canvas, source_weight=source_weight
),
"source": record["source"],
"length": len(canvas.prompt_ids) + len(canvas.canvas_ids),
"record": record,
"max_questions": int(_value(decision, "max_questions_per_canvas", 20)),
}
if slot_plan is not None:
row["slot_plan"] = slot_plan
if _value(latent, "sample_num_slots", False):
row["slot_sampling"] = {
"seed": int(_value(cfg, "seed", 0)),
"max_slots": len(slot_plan.ids),
"pad_token_id": required_token_id(
_value(tokenizer, "pad_token_id"), "tokenizer.pad_token_id"
),
"padding_pinned": True,
}
return row
def _build_canvas_worker(
item: tuple[int, dict[str, Any], float],
*,
tokenizer: Any,
cfg: Any,
spec: Any,
vocab_size: int,
slot_plan: SlotPlan | None,
thought_open_ids: tuple[int, ...],
thought_close_ids: tuple[int, ...],
) -> tuple[int, dict[str, Any] | None, bool]:
index, record, source_weight = item
try:
row = _canvas_row(
tokenizer,
record,
cfg,
source_weight,
spec=spec,
vocab_size=vocab_size,
slot_plan=slot_plan,
thought_open_ids=thought_open_ids,
thought_close_ids=thought_close_ids,
)
except SchemaError as error:
if not _is_canvas_overflow(error):
raise
return index, None, True
return index, row, False
class DecisionDataset(Dataset):
"""In-memory rows retaining typed canvases and decision-label examples."""
def __init__(
self,
rows: Sequence[dict[str, Any]],
manifest: Mapping[str, Any],
*,
expose_lengths: bool = True,
):
self._rows = tuple(rows)
self.manifest = dict(manifest)
self.column_names = ("length",) if expose_lengths else ()
def __len__(self) -> int:
return len(self._rows)
def __getitem__(
self, index: int | str | DecisionDraw | Sequence[int]
) -> dict[str, Any] | list[int] | list[dict[str, Any]]:
if isinstance(index, str):
if index != "length":
raise KeyError(index)
return [int(row["length"]) for row in self._rows]
if isinstance(index, int):
return self._rows[index]
if isinstance(index, Sequence):
rows = []
for item in index:
if not isinstance(item, int):
raise TypeError("packed decision indices must be integers")
rows.append(self._rows[item])
return rows
maximum = self._rows[index.index]
recipe = maximum.get("slot_sampling")
plan = maximum.get("slot_plan")
if not isinstance(recipe, Mapping) or not isinstance(plan, SlotPlan):
raise ValueError("decision draw requires a sampled maximum slot row")
count = sample_slot_count(
seed=int(recipe["seed"]),
epoch=index.epoch,
global_draw_ordinal=index.global_draw_ordinal,
max_slots=int(recipe["max_slots"]),
)
canvas = project_max_canvas(
maximum["canvas"],
plan,
count,
pad_token_id=int(recipe["pad_token_id"]),
padding_pinned=bool(recipe["padding_pinned"]),
)
row = dict(maximum)
row["canvas"] = canvas
row["slot_plan"] = project_slot_plan(plan, count)
row["decision_example"] = decision_example_from_canvas(
canvas, source_weight=maximum["decision_example"].source_weight
)
row["decision_slot_count"] = count
row["decision_draw"] = index
return row
def remove_columns(self, columns: Sequence[str]) -> "DecisionDataset":
if tuple(columns) != ("length",) or "length" not in self.column_names:
raise ValueError("DecisionDataset only exposes the derived length column")
return DecisionDataset(self._rows, self.manifest, expose_lengths=False)
def _is_canvas_overflow(error: SchemaError) -> bool:
message = str(error)
return (
message.startswith("answer template is ") and " canvas holds " in message
) or (
message.startswith("question ")
and " alone needs " in message
and message.endswith(" canvas rows")
)
def _canvas_rows(
tokenizer: Any,
records: Sequence[dict[str, Any]],
cfg: Any,
weights: Mapping[str, float],
spec,
) -> tuple[list[dict[str, Any]], int]:
rows: list[dict[str, Any]] = []
canvas_too_long = 0
vocab_size = _value(model_config_overrides(cfg), "vocab_size")
if vocab_size is None:
vocab_size = len(tokenizer)
slot_plan, thought_open_ids, thought_close_ids = _resolve_slot_plan(
tokenizer, cfg, spec, int(vocab_size)
)
workers = _value(cfg, "dataset_num_proc", 1) or 1
workers = int(workers)
if workers > 1:
LOG.info(
"Preparing %s decision canvases with %s worker threads",
len(records),
workers,
)
work = (
(index, record, float(weights.get(record["source"], 1.0)))
for index, record in enumerate(records)
)
worker = partial(
_build_canvas_worker,
tokenizer=tokenizer,
cfg=cfg,
spec=spec,
vocab_size=int(vocab_size),
slot_plan=slot_plan,
thought_open_ids=tuple(thought_open_ids),
thought_close_ids=tuple(thought_close_ids),
)
with ThreadPoolExecutor(max_workers=workers) as executor:
prepared = list(executor.map(worker, work))
prepared.sort(key=lambda item: item[0])
rows = [row for _, row, _ in prepared if row is not None]
return rows, sum(overflow for _, _, overflow in prepared)
for index, record in enumerate(records):
if index % 1000 == 0:
LOG.info("Preparing decision canvases: %s/%s", index, len(records))
try:
rows.append(
_canvas_row(
tokenizer,
record,
cfg,
float(weights.get(record["source"], 1.0)),
spec=spec,
vocab_size=int(vocab_size),
slot_plan=slot_plan,
thought_open_ids=thought_open_ids,
thought_close_ids=thought_close_ids,
)
)
except SchemaError as error:
if not _is_canvas_overflow(error):
raise
canvas_too_long += 1
return rows, canvas_too_long
def _filter_budget_rows(
rows: Sequence[dict[str, Any]], cfg: Any, spec, *, is_eval: bool = False
) -> tuple[list[dict[str, Any]], dict[str, int]]:
logical_limit = _value(cfg, "sequence_len")
sample_packed = bool(_value(cfg, "sample_packing", False)) and (
not is_eval or _value(cfg, "eval_sample_packing") is not False
)
packed = bool(_value(cfg, "batch_flattening", False)) or sample_packed
resolved = (
resolve_native_packing_budget(
DictDefault(cfg) if isinstance(cfg, Mapping) else cfg,
packed=packed,
batch_size=(
_value(cfg, "micro_batch_size")
if sample_packed
else _value(cfg, "eval_batch_size")
if is_eval
else _value(cfg, "micro_batch_size")
),
)
if packed
else None
)
physical_limit = None if resolved is None else resolved.payload_capacity
kept: list[dict[str, Any]] = []
drops = {"logical": 0, "physical": 0}
for row in rows:
canvas = row["canvas"]
reason = exceeds_budget(
decision_budget(
prompt_tokens=len(canvas.prompt_ids),
canvas_tokens=len(canvas.canvas_ids),
layout=spec.layout,
),
logical_limit=None if logical_limit is None else int(logical_limit),
physical_limit=None if physical_limit is None else int(physical_limit),
)
if reason is None:
kept.append(row)
else:
drops[reason] += 1
return kept, drops
def load_decision_datasets(
cfg: Any, preprocess: bool = False, *, tokenizer: Any = None
) -> TrainDatasetMeta:
"""Load, normalize, decontaminate, canvasize, and return plugin-owned splits."""
del preprocess
decision = _value(cfg, "diffusion_decision")
spec = require_diffusion_spec(cfg)
train_entries = [
entry
for entry in (_value(cfg, "datasets", []) or [])
if str(_value(entry, "type", "")).startswith("diffusion_decision.")
]
test_entries = [
entry
for entry in (_value(cfg, "test_datasets", []) or [])
if str(_value(entry, "type", "")).startswith("diffusion_decision.")
]
if not train_entries:
raise ValueError("diffusion_decision requires at least one decision dataset")
if tokenizer is None:
tokenizer = load_tokenizer(cfg)
prepared_path = _value(cfg, "dataset_prepared_path")
cache_identity = None
if (
isinstance(prepared_path, (str, Path))
and str(prepared_path)
and _prepared_cache_eligible(cfg)
):
source_paths = _prepared_cache_source_paths([*train_entries, *test_entries])
if source_paths is not None:
cache_identity = prepared_cache_identity(cfg, tokenizer, spec, source_paths)
else:
LOG.info("Decision prepared-row cache bypassed: nonlocal or opaque source")
elif isinstance(prepared_path, (str, Path)) and str(prepared_path):
LOG.info("Decision prepared-row cache bypassed: unsupported latent mode")
if cache_identity is not None:
cache_key, cache_payload = cache_identity
cached = load_prepared_cache(Path(prepared_path), cache_key, cache_payload)
if cached is not None:
cached_train, cached_eval, cached_manifest = cached
LOG.info("Decision prepared-row cache hit")
return TrainDatasetMeta(
train_dataset=DecisionDataset(cached_train, cached_manifest),
eval_dataset=DecisionDataset(cached_eval, cached_manifest)
if cached_eval
else None,
total_num_steps=None,
)
LOG.info("Decision prepared-row cache miss")
train: list[dict[str, Any]] = []
explicit_eval: list[dict[str, Any]] = []
protected_eval: list[dict[str, Any]] = []
for entry, force_dev in [
*((entry, False) for entry in train_entries),
*((entry, True) for entry in test_entries),
]:
adapter = str(_value(entry, "type"))
LOG.info(
"Loading decision source %s (%s)",
_value(entry, "path"),
_value(entry, "split", "train"),
)
normalized = [
normalize_record(
adapter,
row,
training=not (_is_eval(entry) or force_dev),
source_split=str(_value(entry, "split", "train")).lower(),
codebook=_value(_value(decision, "labels"), "codebook", "vendored26"),
)
for row in _rows(entry)
]
split = str(_value(entry, "split", "train")).lower()
if split in {"test", "calibration", "ood"}:
protected_eval.extend(normalized)
elif force_dev or split in {"dev", "validation", "eval"}:
explicit_eval.extend(normalized)
elif _is_eval(entry):
protected_eval.extend(normalized)
else:
train.extend(normalized)
if explicit_eval:
dev = explicit_eval
else:
train, dev = _family_dev_split(train)
LOG.info(
"Decontaminating %s training records against %s development and %s protected records",
len(train),
len(dev),
len(protected_eval),
)
mixture = _value(decision, "mixture")
premixed = bool(_value(mixture, "premixed", False))
retained, drops = decontaminate(train, [*dev, *protected_eval])
if premixed and len(retained) != len(train):
raise ValueError("premixed decision rows overlap an evaluation split")
train = [dict(record) for record in retained]
max_questions = int(_value(decision, "max_questions_per_canvas", 20))
if premixed:
_validate_premixed_records(train)
else:
LOG.info("Grouping compatible decision records")
train = group_records(train, max_questions=max_questions)
dev = group_records(dev, max_questions=max_questions)
assert_split_isolation({"train": train, "eval": dev})
weights = _value(_value(decision, "mixture"), "loss_weight", {}) or {}
train_canvas_too_long = 0
train_budget_drops = {"logical": 0, "physical": 0}
pools = None
sampled_batches: tuple[tuple[MixtureSample[dict[str, Any]], ...], ...] = ()
sampler = None
if premixed:
train_rows, train_canvas_too_long = _canvas_rows(
tokenizer, [dict(record) for record in train], cfg, weights, spec
)
train_rows, train_budget_drops = _filter_budget_rows(train_rows, cfg, spec)
if len(train_rows) != len(train):
raise ValueError(
"premixed decision rows must all pass canvas and budget validation"
)
else:
pools = prepare_source_pools(
train,
(),
max_examples_per_source=_value(mixture, "max_examples_per_source"),
seed=int(_value(cfg, "seed", 0)),
)
usable_pools: dict[str, tuple[dict[str, Any], ...]] = {}
for source, records in pools.records.items():
rows, dropped = _canvas_rows(
tokenizer, [dict(record) for record in records], cfg, weights, spec
)
train_canvas_too_long += dropped
rows, budget_drops = _filter_budget_rows(rows, cfg, spec)
for reason, count in budget_drops.items():
train_budget_drops[reason] += count
if rows:
usable_pools[source] = tuple(rows)
configured_weights = _value(mixture, "weights")
if isinstance(configured_weights, Mapping):
configured_weights = {
source: configured_weights[source]
for source in usable_pools
if source in configured_weights
}
sampler = DeterministicMixtureSampler(
usable_pools,
weights=configured_weights,
temperature=_value(mixture, "temperature"),
seed=int(_value(cfg, "seed", 0)),
)
batch_size = int(_value(cfg, "micro_batch_size", 1))
sampled_batches = tuple(
sampler.epoch_batches(
batch_size,
stratified=bool(_value(mixture, "per_batch_stratified", True)),
)
)
train_rows = [
dict(sample.record) for batch in sampled_batches for sample in batch
]
eval_rows, eval_canvas_too_long = _canvas_rows(tokenizer, dev, cfg, weights, spec)
eval_rows, eval_budget_drops = _filter_budget_rows(
eval_rows, cfg, spec, is_eval=True
)
drop_counts = dict(drops)
if pools is not None:
drop_counts.update(
{
key: value
for key, value in pools.counts.items()
if key not in drop_counts
}
)
manifest = {
"train_rows": len(train_rows),
"eval_rows": len(eval_rows),
"protected_eval_rows": len(protected_eval),
"train_sources": sorted({row["source"] for row in train}),
"eval_sources": sorted({row["source"] for row in dev}),
"dropped": drop_counts,
"canvas_too_long": {
"train": train_canvas_too_long,
"eval": eval_canvas_too_long,
},
"budget_drops": {"train": train_budget_drops, "eval": eval_budget_drops},
"mixture_probabilities": dict(sampler.probabilities) if sampler else {},
"premixed": premixed,
"per_batch_stratified": bool(_value(mixture, "per_batch_stratified", True)),
"stratified_micro_batch_size": int(_value(cfg, "micro_batch_size", 1)),
"mixture_seed": int(_value(cfg, "seed", 0)),
"stratified_epoch_batches": tuple(
tuple(sample.source for sample in batch) for batch in sampled_batches
),
}
if any(
_value(record.get("source_metadata", {}), "official_split") is not None
for record in [*train, *dev, *protected_eval]
if record["source"] == "procedural"
):
manifest["procedural_grouping_limit"] = (
"official partitions and state decontamination only; no parent-family key"
)
if cache_identity is not None:
store_prepared_cache(
Path(prepared_path),
cache_key,
cache_payload,
cfg,
train_rows,
eval_rows,
manifest,
)
LOG.info(
"Prepared %s training draws and %s evaluation rows; budget drops: %s",
len(train_rows),
len(eval_rows),
manifest["budget_drops"],
)
return TrainDatasetMeta(
train_dataset=DecisionDataset(train_rows, manifest),
eval_dataset=DecisionDataset(eval_rows, manifest) if eval_rows else None,
total_num_steps=None,
)
src/axolotl/integrations/diffusion_decision/grouping.py
+150
View File
@@ -0,0 +1,150 @@
"""Group compatible normalized decision records without changing their labels."""
from __future__ import annotations
import copy
import json
from collections import OrderedDict
from collections.abc import Mapping, Sequence
from typing import Any
def group_records(
records: Sequence[Mapping[str, Any]], *, max_questions: int
) -> list[dict[str, Any]]:
"""Combine only records with identical serving context and split without loss."""
if max_questions <= 0:
raise ValueError("max_questions must be positive")
groups: OrderedDict[tuple[str, str, str, bytes, bytes], list[Mapping[str, Any]]] = (
OrderedDict()
)
for record in records:
groups.setdefault(_key(record), []).append(record)
result: list[dict[str, Any]] = []
for group in groups.values():
result.extend(_split_group(group, max_questions))
return result
def _key(record: Mapping[str, Any]) -> tuple[str, str, str, bytes, bytes]:
source = record.get("source")
group = record.get("group")
if not isinstance(source, str) or not source:
raise ValueError("decision record requires a nonempty source")
if not isinstance(group, str) or not group:
raise ValueError("decision record requires a nonempty source group")
family = record.get("family", group)
if not isinstance(family, str) or not family:
raise ValueError("decision record requires a nonempty family when present")
return (
source,
group,
family,
_bytes(record.get("state")),
_bytes(record.get("instructions", "")),
)
def _bytes(value: Any) -> bytes:
if isinstance(value, str):
return value.encode("utf-8")
return json.dumps(
value,
ensure_ascii=False,
separators=(",", ":"),
sort_keys=False,
allow_nan=False,
).encode("utf-8")
def _split_group(
records: Sequence[Mapping[str, Any]], max_questions: int
) -> list[dict[str, Any]]:
if len(records) == 1:
record = records[0]
questions = record.get("questions")
labels = record.get("labels")
if not isinstance(questions, Mapping) or not isinstance(labels, Mapping):
raise ValueError("decision record requires questions and labels mappings")
if not questions:
raise ValueError("decision record requires at least one question")
if set(questions) != set(labels):
raise ValueError("decision record question and label ids differ")
if len(questions) <= max_questions:
return [copy.deepcopy(dict(record))]
chunks: list[dict[str, Any]] = []
current: dict[str, Any] | None = None
used: set[str] = set()
seen_ids: set[str] = set()
for position, record in enumerate(records):
questions = record.get("questions")
labels = record.get("labels")
record_id = _record_id(record, position)
if not isinstance(questions, Mapping) or not isinstance(labels, Mapping):
raise ValueError("decision record requires questions and labels mappings")
if not questions:
raise ValueError("decision record requires at least one question")
if record_id in seen_ids:
raise ValueError(f"decision record id is duplicated: {record_id}")
seen_ids.add(record_id)
if set(questions) != set(labels):
raise ValueError("decision record question and label ids differ")
for question_id, question in questions.items():
if not isinstance(question_id, str) or not question_id:
raise ValueError("decision question ids must be nonempty strings")
if current is None or len(current["questions"]) == max_questions:
current = _new_chunk(records[0], len(chunks))
chunks.append(current)
used = set()
output_id = _unique_id(question_id, used)
used.add(output_id)
current["questions"][output_id] = copy.deepcopy(question)
current["labels"][output_id] = copy.deepcopy(labels[question_id])
current["grouped_record_ids"].append(record_id)
current["grouped_source_metadata"][record_id] = copy.deepcopy(
record.get("source_metadata", {})
)
current["grouped_question_metadata"][output_id] = {
"record_id": record_id,
"original_question_id": question_id,
"source_metadata": copy.deepcopy(record.get("source_metadata", {})),
}
for chunk in chunks:
chunk["grouped_record_ids"] = tuple(dict.fromkeys(chunk["grouped_record_ids"]))
return chunks
def _record_id(record: Mapping[str, Any], position: int) -> str:
record_id = record.get("id")
if isinstance(record_id, str) and record_id:
return record_id
if record_id is not None:
raise ValueError("decision record id must be a nonempty string")
return f"{record.get('group')}#record{position}"
def _new_chunk(record: Mapping[str, Any], index: int) -> dict[str, Any]:
result = {
key: copy.deepcopy(value)
for key, value in record.items()
if key not in {"questions", "labels", "id", "source_metadata"}
}
result["id"] = f"{_record_id(record, 0)}#group{index}"
result["source_metadata"] = copy.deepcopy(record.get("source_metadata", {}))
result["questions"] = OrderedDict()
result["labels"] = OrderedDict()
result["grouped_record_ids"] = []
result["grouped_question_metadata"] = {}
result["grouped_source_metadata"] = {}
return result
def _unique_id(question_id: str, used: set[str]) -> str:
if question_id not in used:
return question_id
candidate = f"{question_id}__2"
suffix = 2
while candidate in used:
suffix += 1
candidate = f"{question_id}__{suffix}"
return candidate
src/axolotl/integrations/diffusion_decision/hygiene.py
+86
View File
@@ -0,0 +1,86 @@
"""State decontamination and family boundaries for decision datasets."""
from __future__ import annotations
import hashlib
import json
import unicodedata
from collections.abc import Iterable, Mapping
from typing import Any
from ._util import canonical_json
def state_fingerprint(state: Any) -> str:
"""Hash canonical state content independently of questions and targets."""
if state is None:
raise ValueError("decision record requires a state")
if isinstance(state, str):
try:
state = json.loads(state)
except json.JSONDecodeError:
pass
if isinstance(state, str):
text = " ".join(unicodedata.normalize("NFC", state).split())
else:
text = canonical_json(state, allow_nan=False)
return hashlib.sha256(text.encode("utf-8")).hexdigest()
def family_key(record: Mapping[str, Any]) -> tuple[str, str]:
source = record.get("source")
family = record.get("family") or record.get("group")
if not isinstance(source, str) or not source:
raise ValueError("decision record requires a source")
if not isinstance(family, str) or not family:
raise ValueError("decision record requires a family or group")
return source, family
def decontaminate(
training: Iterable[Mapping[str, Any]],
evaluation: Iterable[Mapping[str, Any]],
*,
exclude_families: bool = True,
) -> tuple[list[Mapping[str, Any]], dict[str, int]]:
"""Remove evaluation states across sources and optionally held-out families."""
eval_states: set[str] = set()
eval_families: set[tuple[str, str]] = set()
for record in evaluation:
eval_states.add(state_fingerprint(record["state"]))
eval_families.add(family_key(record))
kept = []
counts = {"input": 0, "state_overlap": 0, "family_overlap": 0, "kept": 0}
for record in training:
counts["input"] += 1
key = family_key(record)
if state_fingerprint(record["state"]) in eval_states:
counts["state_overlap"] += 1
elif exclude_families and key in eval_families:
counts["family_overlap"] += 1
else:
kept.append(record)
counts["kept"] += 1
return kept, counts
def assert_split_isolation(
splits: Mapping[str, Iterable[Mapping[str, Any]]],
) -> None:
"""Reject state or source-family leakage between named dataset splits."""
states: dict[str, str] = {}
families: dict[tuple[str, str], str] = {}
for split, records in splits.items():
for record in records:
fingerprint = state_fingerprint(record["state"])
key = family_key(record)
if fingerprint in states and states[fingerprint] != split:
raise ValueError(
f"state overlaps splits {states[fingerprint]!r} and {split!r}"
)
if key in families and families[key] != split:
raise ValueError(
f"family {key!r} overlaps splits {families[key]!r} and {split!r}"
)
states[fingerprint] = split
families[key] = split
src/axolotl/integrations/diffusion_decision/loss.py
+926
View File
@@ -0,0 +1,926 @@
"""Typed label objectives for decision canvases."""
from __future__ import annotations
import math
from collections.abc import Callable
from dataclasses import dataclass
from typing import Literal, Mapping, Sequence, TypeAlias
import torch
import torch.nn.functional as F
from .records import DecisionCanvas
LabelSoftmax = Literal["restricted", "full", "both"]
FullCEWeighting = Literal["ce", "dft"]
@dataclass(frozen=True)
class HardLabel:
"""A single correct index among a question's allowed token IDs."""
gold_index: int
smoothing: float | None = None
@dataclass(frozen=True)
class DistributionLabel:
"""A dense or sparse probability target for a question's allowed token IDs.
``candidate_indices`` makes ``probabilities`` sparse. ``other_probability``
is the target mass for every vocabulary token outside those candidates and is
only meaningful with the full-vocabulary objective.
"""
probabilities: tuple[float, ...]
candidate_indices: tuple[int, ...] | None = None
other_probability: float = 0.0
candidate_ids: tuple[str, ...] | None = None
@dataclass(frozen=True)
class SetLabel:
"""Indices of every correct alternative among the allowed token IDs."""
allowed_indices: tuple[int, ...]
DecisionLabelTarget: TypeAlias = HardLabel | DistributionLabel | SetLabel
@dataclass(frozen=True)
class DecisionLabelQuestion:
"""One supervised canvas position and its typed label target."""
position: int
allowed_token_ids: tuple[int, ...]
target: DecisionLabelTarget
weight: float = 1.0
@dataclass(frozen=True)
class DecisionLabelExample:
"""Questions from one canvas, weighted as one training example."""
questions: tuple[DecisionLabelQuestion, ...]
source_weight: float = 1.0
@dataclass(frozen=True)
class DecisionLossResult:
"""Source-weighted question/example averages for logging and optimization."""
loss: torch.Tensor
restricted_loss: torch.Tensor
full_vocab_loss: torch.Tensor
brier_loss: torch.Tensor
effective_full_vocab_loss: torch.Tensor | None = None
full_vocab_dft_hard_weight_sum: torch.Tensor | None = None
full_vocab_dft_hard_count: torch.Tensor | None = None
per_example_loss: tuple[torch.Tensor, ...] | None = None
per_example_restricted_loss: tuple[torch.Tensor, ...] | None = None
per_example_full_vocab_loss: tuple[torch.Tensor, ...] | None = None
per_example_effective_full_vocab_loss: tuple[torch.Tensor, ...] | None = None
per_example_full_vocab_dft_hard_weight_sum: tuple[torch.Tensor, ...] | None = None
per_example_full_vocab_dft_hard_count: tuple[torch.Tensor, ...] | None = None
per_example_brier_loss: tuple[torch.Tensor, ...] | None = None
def _optional_smoothing(value: object) -> float | None:
if value is None:
return None
if (
isinstance(value, bool)
or not isinstance(value, (float, int))
or not math.isfinite(value)
or not 0.0 <= value < 1.0
):
raise ValueError("hard target smoothing must be finite and in [0, 1)")
return float(value)
def label_target_from_mapping(target: Mapping[str, object]) -> DecisionLabelTarget:
"""Convert a normalized-record target into the typed loss contract."""
kind = target.get("kind")
if kind == "hard":
return HardLabel(
_index(target.get("gold_idx"), "gold_idx"),
_optional_smoothing(target.get("smoothing")),
)
if kind == "dist":
values = target.get("probs")
if not isinstance(values, (list, tuple)) or not values:
raise ValueError("dist target requires a nonempty probs sequence")
probabilities = tuple(_probability(value) for value in values)
indices = target.get("candidate_indices")
if indices is None:
candidate_indices = None
else:
if not isinstance(indices, (list, tuple)) or len(indices) != len(values):
raise ValueError("dist candidate_indices must align with probs")
candidate_indices = tuple(
_index(value, "candidate_indices") for value in indices
)
other_probability = _probability(target.get("other_probability", 0.0))
candidate_ids_value = target.get("candidate_ids")
if candidate_ids_value is None:
candidate_ids = None
elif (
not isinstance(candidate_ids_value, (list, tuple))
or len(candidate_ids_value) != len(probabilities)
or any(
not isinstance(value, str) or not value for value in candidate_ids_value
)
):
raise ValueError("dist candidate_ids must align with probs")
else:
candidate_ids = tuple(candidate_ids_value)
_validate_distribution_values(
probabilities, candidate_indices, other_probability
)
return DistributionLabel(
probabilities, candidate_indices, other_probability, candidate_ids
)
if kind == "set":
values = target.get("allowed_set")
if not isinstance(values, (list, tuple)) or not values:
raise ValueError("set target requires a nonempty allowed_set sequence")
indices = tuple(_index(value, "allowed_set") for value in values)
if len(set(indices)) != len(indices):
raise ValueError("set target indices must be unique")
return SetLabel(indices)
raise ValueError(f"unknown decision label kind {kind!r}")
def decision_example_from_canvas(
canvas: DecisionCanvas, *, source_weight: float = 1.0
) -> DecisionLabelExample:
"""Adapt a validated decision canvas for the label-loss API."""
questions = tuple(
DecisionLabelQuestion(
position=int(position),
allowed_token_ids=tuple(int(token) for token in allowed_ids),
target=(
target
if isinstance(target, (HardLabel, DistributionLabel, SetLabel))
else label_target_from_mapping(target)
),
weight=_question_weight(target),
)
for position, allowed_ids, target in zip(
canvas.label_positions, canvas.allowed_ids, canvas.targets, strict=True
)
)
return DecisionLabelExample(questions=questions, source_weight=source_weight)
def decision_label_loss(
logits: torch.Tensor,
examples: Sequence[DecisionLabelExample],
supervision_mask: torch.Tensor,
*,
label_softmax: LabelSoftmax = "both",
full_ce_weighting: FullCEWeighting = "ce",
brier_weight: float = 0.1,
hard_label_smoothing: float = 0.0,
) -> DecisionLossResult:
"""Compute typed label losses using only supervised label positions.
``both`` applies restricted and full-vocabulary CE to every label kind.
Set labels use a detached highest-scoring member of the correct set for
full-vocabulary CE while retaining their set-mass objective.
"""
_validate_inputs(
logits,
examples,
supervision_mask,
label_softmax,
full_ce_weighting,
brier_weight,
hard_label_smoothing,
)
restricted_examples = []
full_vocab_examples = []
effective_full_vocab_examples = []
full_vocab_dft_hard_weight_sum_examples = []
full_vocab_dft_hard_count_examples = []
brier_examples = []
for batch_index, example in enumerate(examples):
_validate_source_weight(example.source_weight)
restricted_questions = []
full_vocab_questions = []
effective_full_vocab_questions = []
full_vocab_dft_hard_weight_sums = []
full_vocab_dft_hard_counts = []
brier_questions = []
for question in example.questions:
query_logits = logits[batch_index, question.position].float()
allowed_ids = _allowed_ids(
question.allowed_token_ids, query_logits.shape[0]
)
allowed_ids = allowed_ids.to(query_logits.device)
restricted_logprobs = F.log_softmax(
query_logits.index_select(0, allowed_ids), dim=0
)
restricted_probs = restricted_logprobs.exp()
if (
isinstance(question.target, DistributionLabel)
and question.target.other_probability > 0
and label_softmax == "full"
):
restricted = query_logits.new_zeros(())
brier = query_logits.new_zeros(())
full_target = allowed_ids[0]
full_distribution = None
else:
restricted, brier, full_target, full_distribution = (
_restricted_question_loss(
restricted_logprobs,
restricted_probs,
allowed_ids,
question.target,
hard_label_smoothing=hard_label_smoothing,
)
)
zero = query_logits.new_zeros(())
if label_softmax == "restricted":
full_vocab = zero
effective_full_vocab = zero
dft_hard_weight_sum = zero
dft_hard_count = zero
else:
full_logprobs = F.log_softmax(query_logits, dim=0)
full_vocab = _full_vocab_question_loss(
full_logprobs,
allowed_ids,
question.target,
full_target,
full_distribution,
)
dft_hard = full_ce_weighting == "dft" and isinstance(
question.target, HardLabel
)
dft_weight = (-full_vocab.detach()).exp() if dft_hard else zero
effective_full_vocab = (
full_vocab * dft_weight if dft_hard else full_vocab
)
dft_hard_weight_sum = dft_weight
dft_hard_count = zero.new_ones(()) if dft_hard else zero
if label_softmax == "full":
restricted = zero
restricted_questions.append(restricted)
full_vocab_questions.append(full_vocab)
effective_full_vocab_questions.append(effective_full_vocab)
full_vocab_dft_hard_weight_sums.append(dft_hard_weight_sum)
full_vocab_dft_hard_counts.append(dft_hard_count)
brier_questions.append(brier)
weight = example.source_weight
qw = [question.weight for question in example.questions]
restricted_examples.append(_weighted_mean(restricted_questions, qw) * weight)
full_vocab_examples.append(_weighted_mean(full_vocab_questions, qw) * weight)
effective_full_vocab_examples.append(
_weighted_mean(effective_full_vocab_questions, qw) * weight
)
full_vocab_dft_hard_weight_sum_examples.append(
torch.stack(full_vocab_dft_hard_weight_sums).sum()
)
full_vocab_dft_hard_count_examples.append(
torch.stack(full_vocab_dft_hard_counts).sum()
)
brier_examples.append(_weighted_mean(brier_questions, qw) * weight)
restricted_loss = torch.stack(restricted_examples).mean()
full_vocab_loss = torch.stack(full_vocab_examples).mean()
effective_full_vocab_loss = torch.stack(effective_full_vocab_examples).mean()
full_vocab_dft_hard_weight_sum = torch.stack(
full_vocab_dft_hard_weight_sum_examples
).sum()
full_vocab_dft_hard_count = torch.stack(full_vocab_dft_hard_count_examples).sum()
brier_loss = torch.stack(brier_examples).mean()
metric_restricted = (
restricted_examples
if label_softmax != "full"
else [restricted_loss.new_zeros(()) for _ in examples]
)
return DecisionLossResult(
loss=restricted_loss + effective_full_vocab_loss + brier_weight * brier_loss,
restricted_loss=restricted_loss,
full_vocab_loss=full_vocab_loss,
brier_loss=brier_loss,
effective_full_vocab_loss=effective_full_vocab_loss,
full_vocab_dft_hard_weight_sum=full_vocab_dft_hard_weight_sum,
full_vocab_dft_hard_count=full_vocab_dft_hard_count,
per_example_loss=tuple(
restricted.detach()
+ effective_full.detach()
+ brier_weight * brier.detach()
for restricted, effective_full, brier in zip(
metric_restricted,
effective_full_vocab_examples,
brier_examples,
strict=True,
)
),
per_example_restricted_loss=(
tuple(value.detach() for value in metric_restricted)
),
per_example_full_vocab_loss=tuple(
value.detach() for value in full_vocab_examples
),
per_example_effective_full_vocab_loss=tuple(
value.detach() for value in effective_full_vocab_examples
),
per_example_full_vocab_dft_hard_weight_sum=tuple(
value.detach() for value in full_vocab_dft_hard_weight_sum_examples
),
per_example_full_vocab_dft_hard_count=tuple(
value.detach() for value in full_vocab_dft_hard_count_examples
),
per_example_brier_loss=tuple(value.detach() for value in brier_examples),
)
def decision_label_loss_from_hidden(
hidden: torch.Tensor,
head: torch.nn.Linear,
examples: Sequence[DecisionLabelExample],
supervision_mask: torch.Tensor,
*,
linear_token_loss: Callable[
[torch.Tensor, torch.nn.Linear, torch.Tensor], torch.Tensor
],
label_softmax: LabelSoftmax = "both",
full_ce_weighting: FullCEWeighting = "ce",
brier_weight: float = 0.1,
hard_label_smoothing: float = 0.0,
) -> DecisionLossResult:
"""Compute typed loss without materializing full-vocabulary logits."""
_validate_hidden_inputs(
hidden,
head,
examples,
supervision_mask,
label_softmax,
full_ce_weighting,
brier_weight,
hard_label_smoothing,
)
restricted_examples: list[torch.Tensor] = []
brier_examples: list[torch.Tensor] = []
full_targets: list[torch.Tensor] = []
full_soft_corrections: list[torch.Tensor | None] = []
selected_hidden: list[torch.Tensor] = []
full_example_indices: list[int] = []
full_question_weights: list[float] = []
full_dft_hard: list[bool] = []
residual_full_losses: list[tuple[int, torch.Tensor, float]] = []
for batch_index, example in enumerate(examples):
_validate_source_weight(example.source_weight)
restricted_questions: list[torch.Tensor] = []
brier_questions: list[torch.Tensor] = []
for question in example.questions:
query_hidden = hidden[batch_index, question.position]
projection_hidden = query_hidden.to(dtype=head.weight.dtype)
allowed_ids = _allowed_ids(question.allowed_token_ids, head.weight.shape[0])
allowed_ids = allowed_ids.to(query_hidden.device)
selected_bias = (
head.bias.index_select(0, allowed_ids)
if head.bias is not None
else None
)
restricted_logits = F.linear(
projection_hidden,
head.weight.index_select(0, allowed_ids),
selected_bias,
).float()
restricted_logprobs = F.log_softmax(restricted_logits, dim=0)
if (
isinstance(question.target, DistributionLabel)
and question.target.other_probability > 0
and label_softmax == "full"
):
restricted = query_hidden.new_zeros(())
brier = query_hidden.new_zeros(())
full_target = allowed_ids[0]
full_distribution = None
else:
restricted, brier, full_target, full_distribution = (
_restricted_question_loss(
restricted_logprobs,
restricted_logprobs.exp(),
allowed_ids,
question.target,
hard_label_smoothing=hard_label_smoothing,
)
)
restricted_questions.append(restricted)
brier_questions.append(brier)
if label_softmax != "restricted":
if (
isinstance(question.target, DistributionLabel)
and question.target.other_probability > 0
):
full_logprobs = F.log_softmax(
F.linear(projection_hidden, head.weight, head.bias).float(),
dim=0,
)
residual_full_losses.append(
(
batch_index,
_full_vocab_question_loss(
full_logprobs,
allowed_ids,
question.target,
full_target,
full_distribution,
),
question.weight,
)
)
continue
selected_hidden.append(projection_hidden)
full_targets.append(full_target)
full_question_weights.append(question.weight)
full_soft_corrections.append(
None
if full_distribution is None
else restricted_logits[full_distribution.argmax()]
- (full_distribution * restricted_logits).sum()
)
full_example_indices.append(batch_index)
full_dft_hard.append(
full_ce_weighting == "dft"
and isinstance(question.target, HardLabel)
)
weight = example.source_weight
qw = [question.weight for question in example.questions]
restricted_examples.append(_weighted_mean(restricted_questions, qw) * weight)
brier_examples.append(_weighted_mean(brier_questions, qw) * weight)
restricted_loss = torch.stack(restricted_examples).mean()
brier_loss = torch.stack(brier_examples).mean()
full_vocab_loss = hidden.new_zeros((), dtype=torch.float32)
effective_full_vocab_loss = hidden.new_zeros((), dtype=torch.float32)
full_vocab_dft_hard_weight_sum = hidden.new_zeros((), dtype=torch.float32)
full_vocab_dft_hard_count = hidden.new_zeros((), dtype=torch.float32)
dft_hard_weight_sum_per_example = [
hidden.new_zeros((), dtype=torch.float32) for _ in examples
]
dft_hard_count_per_example = [
hidden.new_zeros((), dtype=torch.float32) for _ in examples
]
if label_softmax != "restricted":
if selected_hidden:
token_losses = linear_token_loss(
torch.stack(selected_hidden), head, torch.stack(full_targets)
)
if token_losses.ndim != 1 or token_losses.shape[0] != len(selected_hidden):
raise ValueError(
"linear_token_loss must return one loss per selected query"
)
corrections = torch.stack(
[
token_losses[index].new_zeros(())
if value is None
else value.to(dtype=token_losses.dtype)
for index, value in enumerate(full_soft_corrections)
]
)
token_losses = token_losses + corrections
else:
token_losses = hidden.new_empty((0,), dtype=torch.float32)
per_example: list[list[torch.Tensor]] = [[] for _ in examples]
per_example_weights: list[list[float]] = [[] for _ in examples]
for loss, batch_index, question_weight in zip(
token_losses, full_example_indices, full_question_weights, strict=True
):
per_example[batch_index].append(loss)
per_example_weights[batch_index].append(question_weight)
for batch_index, loss, question_weight in residual_full_losses:
per_example[batch_index].append(loss)
per_example_weights[batch_index].append(question_weight)
weighted = [
_weighted_mean(losses, per_example_weights[index])
* examples[index].source_weight
for index, losses in enumerate(per_example)
]
full_vocab_loss = torch.stack(weighted).mean()
if full_ce_weighting == "dft":
dft_weights = (-token_losses.detach()).exp()
dft_hard_mask = torch.tensor(
full_dft_hard, device=token_losses.device, dtype=torch.bool
)
effective_token_losses = torch.where(
dft_hard_mask, token_losses * dft_weights, token_losses
)
else:
dft_weights = torch.zeros_like(token_losses)
effective_token_losses = token_losses
effective_per_example: list[list[torch.Tensor]] = [[] for _ in examples]
for loss, dft_weight, dft_hard, batch_index in zip(
effective_token_losses,
dft_weights,
full_dft_hard,
full_example_indices,
strict=True,
):
effective_per_example[batch_index].append(loss)
if dft_hard:
dft_hard_weight_sum_per_example[batch_index] = (
dft_hard_weight_sum_per_example[batch_index] + dft_weight
)
dft_hard_count_per_example[batch_index] = (
dft_hard_count_per_example[batch_index] + 1
)
for batch_index, loss, _question_weight_unused in residual_full_losses:
effective_per_example[batch_index].append(loss)
effective_weighted = [
_weighted_mean(losses, per_example_weights[index])
* examples[index].source_weight
for index, losses in enumerate(effective_per_example)
]
effective_full_vocab_loss = torch.stack(effective_weighted).mean()
full_vocab_dft_hard_weight_sum = torch.stack(
dft_hard_weight_sum_per_example
).sum()
full_vocab_dft_hard_count = torch.stack(dft_hard_count_per_example).sum()
if label_softmax == "full":
restricted_loss = restricted_loss.new_zeros(())
metric_restricted = (
restricted_examples
if label_softmax != "full"
else [restricted_loss.new_zeros(()) for _ in examples]
)
metric_full = (
weighted
if label_softmax != "restricted"
else [full_vocab_loss.new_zeros(()) for _ in examples]
)
return DecisionLossResult(
loss=restricted_loss + effective_full_vocab_loss + brier_weight * brier_loss,
restricted_loss=restricted_loss,
full_vocab_loss=full_vocab_loss,
brier_loss=brier_loss,
effective_full_vocab_loss=effective_full_vocab_loss,
full_vocab_dft_hard_weight_sum=full_vocab_dft_hard_weight_sum,
full_vocab_dft_hard_count=full_vocab_dft_hard_count,
per_example_loss=tuple(
restricted.detach()
+ effective_full.detach()
+ brier_weight * brier.detach()
for restricted, effective_full, brier in zip(
metric_restricted,
effective_weighted if label_softmax != "restricted" else metric_full,
brier_examples,
strict=True,
)
),
per_example_restricted_loss=tuple(
value.detach() for value in metric_restricted
),
per_example_full_vocab_loss=tuple(value.detach() for value in metric_full),
per_example_effective_full_vocab_loss=tuple(
value.detach()
for value in (
effective_weighted if label_softmax != "restricted" else metric_full
)
),
per_example_full_vocab_dft_hard_weight_sum=tuple(
value.detach() for value in dft_hard_weight_sum_per_example
)
if label_softmax != "restricted"
else tuple(hidden.new_zeros(()) for _ in examples),
per_example_full_vocab_dft_hard_count=tuple(
value.detach() for value in dft_hard_count_per_example
),
per_example_brier_loss=tuple(value.detach() for value in brier_examples),
)
def _restricted_question_loss(
restricted_logprobs: torch.Tensor,
restricted_probs: torch.Tensor,
allowed_ids: torch.Tensor,
target: DecisionLabelTarget,
*,
hard_label_smoothing: float = 0.0,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
if isinstance(target, HardLabel):
_validate_index(target.gold_index, allowed_ids.shape[0], "gold_index")
smoothing = (
hard_label_smoothing
if target.smoothing is None
else _optional_smoothing(target.smoothing)
)
hard_desired = F.one_hot(
torch.tensor(target.gold_index, device=restricted_logprobs.device),
num_classes=allowed_ids.shape[0],
).to(dtype=restricted_logprobs.dtype)
if smoothing:
desired = (
hard_desired * (1.0 - smoothing) + smoothing / allowed_ids.shape[0]
)
primary = -(desired * restricted_logprobs).sum()
else:
desired = hard_desired
primary = -restricted_logprobs[target.gold_index]
full_target = allowed_ids[target.gold_index]
full_distribution = desired if smoothing else None
elif isinstance(target, DistributionLabel):
indices, probabilities = _distribution_candidates(target, allowed_ids.shape[0])
if target.other_probability > 0:
raise ValueError(
"a sparse dist target with other_probability requires "
"label_softmax=full"
)
desired = torch.tensor(
probabilities,
dtype=restricted_logprobs.dtype,
device=restricted_logprobs.device,
)
if len(indices) != allowed_ids.shape[0]:
dense_desired = torch.zeros_like(restricted_logprobs)
dense_desired[torch.tensor(indices, device=allowed_ids.device)] = desired
desired = dense_desired
ce_desired = desired
if hard_label_smoothing and _is_one_hot_distribution(
target, allowed_ids.shape[0]
):
ce_desired = (
desired * (1.0 - hard_label_smoothing)
+ hard_label_smoothing / allowed_ids.shape[0]
)
positive = ce_desired > 0
primary = torch.where(
positive,
ce_desired * (ce_desired.log() - restricted_logprobs),
torch.zeros_like(ce_desired),
).sum()
full_target = allowed_ids[desired.argmax()]
full_distribution = ce_desired
elif isinstance(target, SetLabel):
_validate_set(target.allowed_indices, allowed_ids.shape[0])
indices = torch.tensor(
target.allowed_indices, dtype=torch.long, device=restricted_logprobs.device
)
set_logprobs = restricted_logprobs.index_select(0, indices)
primary = -torch.logsumexp(set_logprobs, dim=0)
set_mass = set_logprobs.exp().sum()
full_target = allowed_ids[indices[set_logprobs.detach().argmax()]]
brier = (1.0 - set_mass).square()
full_distribution = None
else:
raise TypeError("unsupported decision label target")
if not isinstance(target, SetLabel):
brier_target = hard_desired if isinstance(target, HardLabel) else desired
brier = torch.sum((restricted_probs - brier_target).square())
return primary, brier, full_target, full_distribution
def _full_vocab_question_loss(
full_logprobs: torch.Tensor,
allowed_ids: torch.Tensor,
target: DecisionLabelTarget,
full_target: torch.Tensor,
full_distribution: torch.Tensor | None,
) -> torch.Tensor:
if isinstance(target, DistributionLabel) and target.other_probability > 0:
indices, probabilities = _distribution_candidates(target, allowed_ids.shape[0])
candidate_ids = allowed_ids[torch.tensor(indices, device=allowed_ids.device)]
desired = torch.tensor(
probabilities, dtype=full_logprobs.dtype, device=full_logprobs.device
)
loss = -(desired * full_logprobs.index_select(0, candidate_ids)).sum()
other_logprob = full_logprobs.index_fill(
0, candidate_ids, -torch.inf
).logsumexp(0)
return loss - target.other_probability * other_logprob
if full_distribution is None:
if not isinstance(target, DistributionLabel):
return -full_logprobs[full_target]
indices, probabilities = _distribution_candidates(target, allowed_ids.shape[0])
candidate_ids = allowed_ids[torch.tensor(indices, device=allowed_ids.device)]
desired = torch.tensor(
probabilities, dtype=full_logprobs.dtype, device=full_logprobs.device
)
else:
candidate_ids = allowed_ids
desired = full_distribution
return -(desired * full_logprobs.index_select(0, candidate_ids)).sum()
def _validate_inputs(
logits: torch.Tensor,
examples: Sequence[DecisionLabelExample],
supervision_mask: torch.Tensor,
label_softmax: LabelSoftmax,
full_ce_weighting: FullCEWeighting,
brier_weight: float,
hard_label_smoothing: float,
) -> None:
if logits.ndim != 3:
raise ValueError("logits must be [batch, sequence, vocabulary]")
if supervision_mask.shape != logits.shape[:2]:
raise ValueError("supervision_mask must be [batch, sequence]")
if supervision_mask.dtype is not torch.bool:
raise TypeError("supervision_mask must be bool")
if supervision_mask.device != logits.device:
raise ValueError("supervision_mask and logits must share a device")
if len(examples) != logits.shape[0] or not examples:
raise ValueError("examples must contain one nonempty entry per batch item")
if label_softmax not in {"restricted", "full", "both"}:
raise ValueError("label_softmax must be restricted, full, or both")
if full_ce_weighting not in {"ce", "dft"}:
raise ValueError("full_ce_weighting must be ce or dft")
if (
isinstance(brier_weight, bool)
or not isinstance(brier_weight, (float, int))
or not math.isfinite(brier_weight)
):
raise ValueError("brier_weight must be finite")
if brier_weight < 0:
raise ValueError("brier_weight must be nonnegative")
if (
isinstance(hard_label_smoothing, bool)
or not isinstance(hard_label_smoothing, (float, int))
or not math.isfinite(hard_label_smoothing)
or not 0.0 <= hard_label_smoothing < 1.0
):
raise ValueError("hard_label_smoothing must be finite and in [0, 1)")
for batch_index, example in enumerate(examples):
if not example.questions:
raise ValueError("each example must have at least one question")
positions: set[int] = set()
for question in example.questions:
if (
isinstance(question.position, bool)
or not isinstance(question.position, int)
or not 0 <= question.position < logits.shape[1]
):
raise ValueError("question position is outside the canvas logits")
if question.position in positions:
raise ValueError("each example requires distinct question positions")
positions.add(question.position)
torch._assert_async(
supervision_mask[batch_index, question.position],
"decision loss requires supervised label positions",
)
def _validate_hidden_inputs(
hidden: torch.Tensor,
head: torch.nn.Linear,
examples: Sequence[DecisionLabelExample],
supervision_mask: torch.Tensor,
label_softmax: LabelSoftmax,
full_ce_weighting: FullCEWeighting,
brier_weight: float,
hard_label_smoothing: float,
) -> None:
_validate_inputs(
hidden,
examples,
supervision_mask,
label_softmax,
full_ce_weighting,
brier_weight,
hard_label_smoothing,
)
if not isinstance(head, torch.nn.Linear) or head.weight.ndim != 2:
raise TypeError("head must be a linear vocabulary projection")
if head.weight.shape[1] != hidden.shape[2]:
raise ValueError("head input features must match hidden states")
if head.bias is not None and head.bias.shape != (head.weight.shape[0],):
raise ValueError("head bias must align with vocabulary rows")
def _allowed_ids(token_ids: tuple[int, ...], vocab_size: int) -> torch.Tensor:
if not token_ids:
raise ValueError("a question requires at least one allowed token ID")
if len(set(token_ids)) != len(token_ids):
raise ValueError("allowed token IDs must be unique")
if any(
isinstance(token, bool)
or not isinstance(token, int)
or not 0 <= token < vocab_size
for token in token_ids
):
raise ValueError("allowed token ID is outside the vocabulary")
return torch.tensor(token_ids, dtype=torch.long)
def _question_weight(target: object) -> float:
"""Optional per-question ``weight`` on a normalized label mapping (default 1)."""
if not isinstance(target, Mapping):
return 1.0
weight = target.get("weight", 1.0)
if (
isinstance(weight, bool)
or not isinstance(weight, (int, float))
or not math.isfinite(weight)
or weight <= 0
):
raise ValueError("label weight must be a positive finite number")
return float(weight)
def _weighted_mean(
losses: Sequence[torch.Tensor], weights: Sequence[float]
) -> torch.Tensor:
"""Question average with per-question weights; weight 1 everywhere is the plain mean."""
stacked = torch.stack(list(losses))
scale = torch.tensor(list(weights), device=stacked.device, dtype=stacked.dtype)
return (stacked * scale).sum() / stacked.shape[0]
def _validate_source_weight(weight: float) -> None:
if (
isinstance(weight, bool)
or not isinstance(weight, (float, int))
or not math.isfinite(weight)
or weight <= 0
):
raise ValueError("source_weight must be a positive finite number")
def _validate_index(index: int, count: int, name: str) -> None:
if isinstance(index, bool) or not isinstance(index, int) or not 0 <= index < count:
raise ValueError(f"{name} must identify an allowed label")
def _distribution_candidates(
target: DistributionLabel, count: int
) -> tuple[tuple[int, ...], tuple[float, ...]]:
indices = (
tuple(range(count))
if target.candidate_indices is None
else target.candidate_indices
)
_validate_distribution_values(
target.probabilities, indices, target.other_probability, count
)
return indices, target.probabilities
def _validate_distribution_values(
probabilities: tuple[float, ...],
candidate_indices: tuple[int, ...] | None,
other_probability: float,
count: int | None = None,
) -> None:
if not probabilities:
raise ValueError("dist target requires at least one candidate")
if candidate_indices is None:
if count is not None and len(probabilities) != count:
raise ValueError("dist target probabilities must align with allowed labels")
elif len(candidate_indices) != len(probabilities):
raise ValueError("dist candidate_indices must align with probs")
elif len(set(candidate_indices)) != len(candidate_indices):
raise ValueError("dist candidate_indices must be unique")
elif count is not None:
for index in candidate_indices:
_validate_index(index, count, "dist candidate index")
if any(
isinstance(value, bool) or not math.isfinite(value) or value < 0
for value in probabilities
):
raise ValueError("dist target probabilities must be finite and nonnegative")
if not math.isfinite(other_probability) or other_probability < 0:
raise ValueError("dist other_probability must be finite and nonnegative")
if not math.isclose(
sum(probabilities) + other_probability, 1.0, rel_tol=0.0, abs_tol=1e-6
):
raise ValueError(
"dist target probabilities plus other_probability must sum to one"
)
def _validate_set(indices: tuple[int, ...], count: int) -> None:
if not indices or len(set(indices)) != len(indices):
raise ValueError("set target indices must be nonempty and unique")
for index in indices:
_validate_index(index, count, "set target index")
def _is_one_hot_distribution(target: DistributionLabel, count: int) -> bool:
return (
target.other_probability == 0.0
and (target.candidate_indices is not None or len(target.probabilities) == count)
and sum(value == 1.0 for value in target.probabilities) == 1
and all(value in {0.0, 1.0} for value in target.probabilities)
)
def _index(value: object, name: str) -> int:
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
raise ValueError(f"{name} must be a nonnegative integer")
return value
def _probability(value: object) -> float:
if isinstance(value, bool) or not isinstance(value, (float, int)):
raise ValueError("dist target probabilities must be numeric")
result = float(value)
if not math.isfinite(result) or result < 0:
raise ValueError("dist target probabilities must be finite and nonnegative")
return result
src/axolotl/integrations/diffusion_decision/manifest.py
+341
View File
@@ -0,0 +1,341 @@
"""Serializable decision-training provenance saved with adapters and checkpoints."""
from __future__ import annotations
__ci_config_keys__ = ("diffusion_decision",)
import json
import math
from pathlib import Path
from typing import Any, Literal, get_type_hints
from pydantic import (
BaseModel,
ConfigDict,
Field,
StrictInt,
create_model,
field_validator,
model_validator,
)
from transformers import TrainerCallback
from axolotl.model_support import DiffusionSpec
from ._util import (
CONTROL_TOKEN_FIELDS,
_value,
optional_token_id,
parse_token_ids,
require_diffusion_spec,
resolve_mask_token_id,
)
from .args import DiffusionDecisionConfig
MANIFEST_FILENAME = "diffusion_decision_manifest.json"
MANIFEST_SCHEMA_VERSION = 1
class ModelIdentity(BaseModel):
"""Resolved base-model identity retained by a decision adapter."""
model_config = ConfigDict(extra="forbid", frozen=True)
base_model: str = Field(min_length=1)
requested_revision: str | None = None
resolved_revision: str | None = None
model_config_type: str | None = None
mask_token_id: StrictInt | None = Field(default=None, ge=0)
class TokenizerSpecialIds(BaseModel):
"""Tokenizer IDs required to reconstruct a decision canvas."""
model_config = ConfigDict(extra="forbid", frozen=True)
bos_token_id: StrictInt | None = Field(default=None, ge=0)
eos_token_id: StrictInt | None = Field(default=None, ge=0)
pad_token_id: StrictInt | None = Field(default=None, ge=0)
unk_token_id: StrictInt | None = Field(default=None, ge=0)
mask_token_id: StrictInt | None = Field(default=None, ge=0)
all_special_ids: list[StrictInt] = Field(default_factory=list)
@field_validator("all_special_ids")
@classmethod
def validate_special_ids(cls, value: list[int]) -> list[int]:
if any(token < 0 for token in value):
raise ValueError("tokenizer special IDs must be nonnegative")
return list(dict.fromkeys(value))
_SPEC_FIELD_OVERRIDES: dict[str, Any] = {
"max_canvas": (StrictInt | None, Field(ge=1, default=None)),
"max_context": (StrictInt | None, Field(ge=1, default=None)),
"time_floor": (float, Field(ge=0.0, le=1.0)),
}
class _DiffusionSpecManifestBase(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)
@model_validator(mode="after")
def validate_spec_contract(self) -> "_DiffusionSpecManifestBase":
DiffusionSpec(**self.model_dump())
return self
DiffusionSpecManifest = create_model(
"DiffusionSpecManifest",
__base__=_DiffusionSpecManifestBase,
__doc__="Validated primitive representation of the resolved model contract.",
**{
name: _SPEC_FIELD_OVERRIDES.get(name, (hint, ...))
for name, hint in get_type_hints(DiffusionSpec).items()
},
)
class _SlotConfiguration(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)
mode: Literal["none", "pad", "pinned", "learned", "free", "mask", "prompt"]
num_slots: StrictInt = Field(ge=0)
token_ids: list[StrictInt] = Field(default_factory=list)
free_update_policy: Literal["argmax"] | None = None
free_slot_init_policy: Literal["prepared_canvas_v0", "fresh_runtime_v1"] = (
"prepared_canvas_v0"
)
@field_validator("token_ids")
@classmethod
def validate_token_ids(cls, value: list[int]) -> list[int]:
if any(token < 0 for token in value):
raise ValueError("configured slot token IDs must be nonnegative")
return value
class ConfiguredSlots(_SlotConfiguration):
"""Static slot configuration without runtime count sampling."""
sample_num_slots: Literal[False] = False
class SampledTrainingSlots(BaseModel):
"""Versioned per-logical-draw slot-count policy for a sampled-slot run."""
model_config = ConfigDict(extra="forbid", frozen=True)
policy: Literal["uniform_inclusive_v1"] = "uniform_inclusive_v1"
min_slots: StrictInt = Field(default=0, ge=0)
max_slots: StrictInt = Field(ge=1)
draw_unit: Literal["logical_draw"] = "logical_draw"
seed_derivation: Literal["seed_epoch_global_draw_ordinal_v1"] = (
"seed_epoch_global_draw_ordinal_v1"
)
@model_validator(mode="after")
def validate_bounds(self) -> "SampledTrainingSlots":
if self.min_slots > self.max_slots:
raise ValueError("sampled slot minimum cannot exceed its maximum")
return self
class FixedMaxEvaluationSlots(BaseModel):
"""Fixed evaluation materialization used for a sampled-slot training run."""
model_config = ConfigDict(extra="forbid", frozen=True)
policy: Literal["fixed_configured_max_v1"] = "fixed_configured_max_v1"
count: StrictInt = Field(ge=1)
class SampledConfiguredSlots(_SlotConfiguration):
"""Configured slots with distinct training-sampling and evaluation policies."""
sample_num_slots: Literal[True] = True
sampled_training: SampledTrainingSlots
fixed_max_evaluation: FixedMaxEvaluationSlots
@model_validator(mode="after")
def validate_sampled_contract(self) -> "SampledConfiguredSlots":
if self.mode == "none" or self.num_slots < 1:
raise ValueError("sampled slots require a non-none mode and positive count")
if self.sampled_training.max_slots != self.num_slots:
raise ValueError("sampled slot maximum must equal configured num_slots")
if self.fixed_max_evaluation.count != self.num_slots:
raise ValueError(
"fixed evaluation slot count must equal configured num_slots"
)
return self
_RETIRED_READ_FIELDS = frozenset({"eval_steps", "noise_draws", "reader"})
class NoiseReadProtocol(BaseModel):
"""Configured decision read and diffusion controls."""
model_config = ConfigDict(extra="forbid", frozen=True)
read_fraction: float
carry_enabled: bool
diffusion_steps: StrictInt = Field(ge=1)
t_eps: float | None = Field(default=None, ge=0.0, le=1.0)
time_weighting: str | None = None
objective_reduction: str | None = None
@model_validator(mode="before")
@classmethod
def drop_retired_fields(cls, data: Any) -> Any:
if isinstance(data, dict):
data = {k: v for k, v in data.items() if k not in _RETIRED_READ_FIELDS}
return data
@model_validator(mode="after")
def validate_protocol(self) -> "NoiseReadProtocol":
if (
not math.isfinite(self.read_fraction)
or not 0.0 <= self.read_fraction <= 1.0
):
raise ValueError("read_fraction must be finite and in [0, 1]")
if self.t_eps is not None and not math.isfinite(self.t_eps):
raise ValueError("t_eps must be finite")
return self
class DecisionManifest(BaseModel):
"""Validated, portable decision-adapter provenance."""
model_config = ConfigDict(extra="forbid", frozen=True)
schema_version: Literal[1] = 1
model: ModelIdentity
diffusion_spec: DiffusionSpecManifest
decision_layout: Literal["thought_block", "prompt_slots"]
canvas_width: StrictInt = Field(ge=1)
label_codebook: Literal["vendored26", "expanded52"] = "vendored26"
configured_slots: ConfiguredSlots | SampledConfiguredSlots
noise_read_protocol: NoiseReadProtocol
tokenizer_special_ids: TokenizerSpecialIds
initial_common_adapter_sha256: str | None = None
initial_common_adapter_fingerprint_reason: str | None = None
initial_common_adapter_devices: list[str] = Field(default_factory=list)
initial_common_adapter_dtypes: list[str] = Field(default_factory=list)
initial_common_adapter_tensor_count: StrictInt = Field(default=0, ge=0)
@classmethod
def from_path(cls, path: str | Path) -> "DecisionManifest":
return cls.model_validate_json(Path(path).read_text(encoding="utf-8"))
def write_to(self, directory: str | Path) -> Path:
destination = Path(directory) / MANIFEST_FILENAME
destination.parent.mkdir(parents=True, exist_ok=True)
destination.write_text(
json.dumps(self.model_dump(mode="json"), indent=2, sort_keys=True) + "\n",
encoding="utf-8",
)
return destination
def _decision_config(cfg: Any) -> DiffusionDecisionConfig:
value = _value(cfg, "diffusion_decision")
if isinstance(value, DiffusionDecisionConfig):
return value
if isinstance(value, dict):
return DiffusionDecisionConfig.model_validate(value)
raise ValueError("diffusion_decision manifest requires diffusion_decision settings")
def build_decision_manifest(
cfg: Any,
*,
tokenizer: Any = None,
model: Any = None,
) -> DecisionManifest:
"""Build a validated manifest without changing model or tokenizer state."""
decision = _decision_config(cfg)
diffusion = _value(cfg, "diffusion_lm")
if diffusion is None:
raise ValueError("diffusion_decision manifest requires diffusion_lm settings")
spec = require_diffusion_spec(cfg)
model_config = _value(model, "config")
mask_token_id = resolve_mask_token_id(diffusion, model_config)
canvas_width = _value(diffusion, "canvas_width")
if canvas_width is None:
canvas_width = 128
return DecisionManifest(
model=ModelIdentity(
base_model=str(_value(cfg, "base_model")),
requested_revision=_value(cfg, "revision_of_model"),
resolved_revision=_value(model_config, "_commit_hash"),
model_config_type=_value(cfg, "model_config_type"),
mask_token_id=mask_token_id,
),
diffusion_spec=DiffusionSpecManifest.model_validate(spec.to_dict()),
decision_layout=decision.layout,
canvas_width=canvas_width,
label_codebook=decision.labels.codebook,
configured_slots=(
SampledConfiguredSlots(
mode=decision.latent.mode,
num_slots=decision.latent.num_slots,
token_ids=list(decision.latent.token_ids),
free_update_policy=decision.latent.free_update_policy,
free_slot_init_policy=(
"fresh_runtime_v1"
if decision.latent.mode == "free"
else "prepared_canvas_v0"
),
sampled_training=SampledTrainingSlots(
max_slots=decision.latent.num_slots
),
fixed_max_evaluation=FixedMaxEvaluationSlots(
count=decision.latent.num_slots
),
)
if decision.latent.sample_num_slots
else ConfiguredSlots(
mode=decision.latent.mode,
num_slots=decision.latent.num_slots,
token_ids=list(decision.latent.token_ids),
sample_num_slots=False,
free_update_policy=decision.latent.free_update_policy,
free_slot_init_policy=(
"fresh_runtime_v1"
if decision.latent.mode == "free"
else "prepared_canvas_v0"
),
)
),
noise_read_protocol=NoiseReadProtocol(
read_fraction=decision.read_fraction,
carry_enabled=decision.carry.enabled,
diffusion_steps=_value(diffusion, "num_diffusion_steps", 1),
t_eps=_value(diffusion, "t_eps"),
time_weighting=_value(diffusion, "time_weighting"),
objective_reduction=_value(diffusion, "objective_reduction"),
),
tokenizer_special_ids=TokenizerSpecialIds(
**{
name: optional_token_id(_value(tokenizer, name), name)
for name in CONTROL_TOKEN_FIELDS
},
all_special_ids=list(
parse_token_ids(_value(tokenizer, "all_special_ids"), "all_special_ids")
),
),
)
class DecisionManifestCheckpointCallback(TrainerCallback):
"""Write validated decision provenance with Trainer checkpoint directories."""
def __init__(self, manifest: DecisionManifest) -> None:
self.manifest = manifest
def on_save(self, args, state, control, **_kwargs):
if state.is_world_process_zero:
self.manifest.write_to(
Path(args.output_dir) / f"checkpoint-{state.global_step}"
)
return control
src/axolotl/integrations/diffusion_decision/mixture.py
+311
View File
@@ -0,0 +1,311 @@
"""Deterministic, decontaminated source mixing for decision training."""
from __future__ import annotations
import math
import random
from collections import Counter, defaultdict
from collections.abc import Iterator, Mapping, Sequence
from dataclasses import dataclass
from typing import Any, Generic, TypeVar
from ._util import stable_seed
from .hygiene import decontaminate
T = TypeVar("T")
@dataclass(frozen=True)
class PreparedSourcePools:
"""Decontaminated and capped normalized records partitioned by source."""
records: Mapping[str, tuple[Mapping[str, Any], ...]]
counts: Mapping[str, int]
@dataclass(frozen=True)
class MixtureSample(Generic[T]):
"""A training example with its source retained for metrics and loss weighting."""
source: str
record: T
class _SourceCursor(Generic[T]):
def __init__(self, records: Sequence[T], seed: int) -> None:
self._records = tuple(records)
self._order = list(range(len(records)))
self._random = random.Random(seed) # nosec B311 - Deterministic training order.
self._random.shuffle(self._order)
self._offset = 0
def next(self) -> T:
if self._offset == len(self._order):
self._random.shuffle(self._order)
self._offset = 0
record = self._records[self._order[self._offset]]
self._offset += 1
return record
class DeterministicMixtureSampler(Generic[T]):
"""Draw source-tagged examples and source-stratified microbatches."""
def __init__(
self,
pools: Mapping[str, Sequence[T]],
*,
weights: Mapping[str, float] | None = None,
temperature: float | None = None,
seed: int = 0,
) -> None:
self._pools = _validate_pools(pools)
self.sources = tuple(sorted(self._pools))
self.probabilities = _probabilities(
self._pools, self.sources, weights, temperature
)
self._probability_values = tuple(
self.probabilities[source] for source in self.sources
)
self._seed = seed
self._source_random = random.Random(seed) # nosec B311 - Deterministic draws.
self._cursors = {
source: _SourceCursor(records, stable_seed(seed, source))
for source, records in self._pools.items()
}
self._fractional_credit = {source: 0.0 for source in self.sources}
def draw(self, count: int) -> tuple[MixtureSample[T], ...]:
"""Draw examples by source probability, cycling deterministically within sources."""
_validate_count(count, "count")
choices = self._source_random.choices(
self.sources, weights=self._probability_values, k=count
)
return tuple(
MixtureSample(source, self._cursors[source].next()) for source in choices
)
def batches(
self, batch_size: int, count: int, *, stratified: bool = True
) -> Iterator[tuple[MixtureSample[T], ...]]:
"""Yield fixed-size deterministic microbatches without flattening examples."""
_validate_count(batch_size, "batch_size", positive=True)
_validate_count(count, "count")
for _ in range(count):
if stratified:
sources = self._stratified_sources(batch_size)
self._source_random.shuffle(sources)
yield tuple(
MixtureSample(source, self._cursors[source].next())
for source in sources
)
else:
yield self.draw(batch_size)
def epoch_batches(
self, batch_size: int, *, stratified: bool = True
) -> Iterator[tuple[MixtureSample[T], ...]]:
"""Yield batches until every capped source record is seen once."""
_validate_count(batch_size, "batch_size", positive=True)
epoch_seed = self._source_random.randrange(1 << 63) # nosec B311 - Reproducible seed.
cursors = {
source: _SourceCursor(records, stable_seed(epoch_seed, source))
for source, records in self._pools.items()
}
credit = {source: 0.0 for source in self.sources}
seen = {source: 0 for source in self.sources}
while not all(
seen[source] >= len(self._pools[source]) for source in self.sources
):
if stratified:
sources = self._stratified_sources(batch_size, credit)
random.Random(epoch_seed).shuffle(sources) # nosec B311 - Reproducible order.
else:
sources = self._source_random.choices(
self.sources, weights=self._probability_values, k=batch_size
)
epoch_seed += 1
batch = tuple(
MixtureSample(source, cursors[source].next()) for source in sources
)
for source in sources:
seen[source] += 1
yield batch
def _stratified_sources(
self, batch_size: int, fractional_credit: dict[str, float] | None = None
) -> list[str]:
credit = (
self._fractional_credit if fractional_credit is None else fractional_credit
)
expected = {
source: self.probabilities[source] * batch_size for source in self.sources
}
fixed = {source: math.floor(value) for source, value in expected.items()}
remaining = batch_size - sum(fixed.values())
selected = [source for source in self.sources for _ in range(fixed[source])]
for source in self.sources:
credit[source] += expected[source] - fixed[source]
for _ in range(remaining):
source = max(
enumerate(self.sources),
key=lambda item: (credit[item[1]], -item[0]),
)[1]
selected.append(source)
credit[source] -= 1.0
return selected
def prepare_source_pools(
training: Sequence[Mapping[str, Any]],
evaluation: Sequence[Mapping[str, Any]],
*,
max_examples_per_source: int | Mapping[str, int] | None = None,
exclude_families: bool = True,
seed: int = 0,
) -> PreparedSourcePools:
"""Decontaminate records, then deterministically cap each nonempty source pool."""
retained, counts = decontaminate(
training, evaluation, exclude_families=exclude_families
)
grouped: dict[str, list[Mapping[str, Any]]] = defaultdict(list)
for record in retained:
grouped[record["source"]].append(record)
_validate_caps(max_examples_per_source, grouped)
pools = {
source: _cap_records(
records, _cap_for_source(max_examples_per_source, source), seed
)
for source, records in grouped.items()
}
records = {source: pool for source, pool in pools.items() if pool}
if not records:
raise ValueError("no nonempty source pools remain after decontamination")
capped = sum(len(grouped[source]) - len(pool) for source, pool in records.items())
result_counts = dict(counts)
result_counts["capped"] = capped
result_counts["sources"] = len(records)
return PreparedSourcePools(records=records, counts=result_counts)
def source_counts(samples: Sequence[MixtureSample[object]]) -> Mapping[str, int]:
"""Count source-tagged samples without inspecting or altering their records."""
return dict(Counter(sample.source for sample in samples))
def _validate_pools(pools: Mapping[str, Sequence[T]]) -> dict[str, tuple[T, ...]]:
if not pools:
raise ValueError("mixture requires at least one source pool")
result = {}
for source, records in pools.items():
if not isinstance(source, str) or not source:
raise ValueError("mixture source names must be nonempty strings")
if not records:
raise ValueError(f"mixture source {source!r} is empty")
result[source] = tuple(records)
return result
def _probabilities(
pools: Mapping[str, Sequence[object]],
sources: tuple[str, ...],
weights: Mapping[str, float] | None,
temperature: float | None,
) -> Mapping[str, float]:
if weights is not None:
if temperature is not None:
raise ValueError("specify explicit weights or temperature, not both")
if set(weights) != set(sources):
raise ValueError("explicit mixture weights must name every nonempty source")
values = {
source: _positive_finite(weights[source], "mixture weight")
for source in sources
}
return _normalize_probabilities(values)
else:
resolved_temperature = 0.5 if temperature is None else temperature
exponent = _nonnegative_finite(resolved_temperature, "mixture temperature")
log_values = {
source: exponent * math.log(len(pools[source])) for source in sources
}
if any(not math.isfinite(value) for value in log_values.values()):
raise ValueError("mixture temperature produces nonfinite source weights")
maximum = max(log_values.values())
values = {
source: math.exp(value - maximum) for source, value in log_values.items()
}
return _normalize_probabilities(values)
def _normalize_probabilities(values: Mapping[str, float]) -> Mapping[str, float]:
maximum = max(values.values())
scaled = {source: value / maximum for source, value in values.items()}
total = sum(scaled.values())
probabilities = {source: value / total for source, value in scaled.items()}
if any(not math.isfinite(value) or value <= 0 for value in probabilities.values()):
raise ValueError("mixture source probability underflowed to zero")
return probabilities
def _validate_caps(
cap: int | Mapping[str, int] | None, pools: Mapping[str, Sequence[object]]
) -> None:
if isinstance(cap, Mapping):
unknown = set(cap) - set(pools)
if unknown:
raise ValueError("per-source cap names an unavailable source")
for value in cap.values():
_positive_integer(value, "max_examples_per_source")
elif cap is not None:
_positive_integer(cap, "max_examples_per_source")
def _cap_for_source(cap: int | Mapping[str, int] | None, source: str) -> int | None:
if isinstance(cap, Mapping):
return cap.get(source)
return cap
def _cap_records(
records: Sequence[Mapping[str, Any]], cap: int | None, seed: int
) -> tuple[Mapping[str, Any], ...]:
if cap is None or len(records) <= cap:
return tuple(records)
order = list(range(len(records)))
random.Random(stable_seed(seed, records[0]["source"])).shuffle( # nosec B311
order
)
return tuple(records[index] for index in order[:cap])
def _positive_finite(value: float, name: str) -> float:
if isinstance(value, bool) or not isinstance(value, (int, float)):
raise ValueError(f"{name} must be positive and finite")
result = float(value)
if not math.isfinite(result) or result <= 0:
raise ValueError(f"{name} must be positive and finite")
return result
def _nonnegative_finite(value: float, name: str) -> float:
if isinstance(value, bool) or not isinstance(value, (int, float)):
raise ValueError(f"{name} must be nonnegative and finite")
result = float(value)
if not math.isfinite(result) or result < 0:
raise ValueError(f"{name} must be nonnegative and finite")
return result
def _positive_integer(value: int, name: str) -> None:
if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
raise ValueError(f"{name} must be a positive integer")
def _validate_count(value: int, name: str, *, positive: bool = False) -> None:
if (
isinstance(value, bool)
or not isinstance(value, int)
or (value < 0 or (positive and value == 0))
):
qualifier = "positive" if positive else "nonnegative"
raise ValueError(f"{name} must be a {qualifier} integer")
src/axolotl/integrations/diffusion_decision/permutation.py
+68
View File
@@ -0,0 +1,68 @@
"""Deterministic question and choice ordering with target preservation."""
from __future__ import annotations
import hashlib
import random
from collections.abc import Mapping
from typing import Any
from ._util import canonical_json
from .adapters.jsonl import normalize_jsonl
def permute_record(
record: Mapping[str, Any], *, seed: int, codebook: str = "vendored26"
) -> dict[str, Any]:
result = normalize_jsonl(record, codebook=codebook)
identity = _permutation_identity(result)
keys = sorted(result["questions"])
_generator(seed, identity, "questions").shuffle(keys)
result["questions"] = {key: result["questions"][key] for key in keys}
result["labels"] = {key: result["labels"][key] for key in keys}
for key, question in result["questions"].items():
if question["type"] != "choice":
continue
options = question["options"]
order = list(range(len(options)))
_generator(seed, identity, f"choices:{key}").shuffle(order)
question["options"] = [options[index] for index in order]
target = result["labels"][key]
inverse = {old: new for new, old in enumerate(order)}
if target["kind"] == "hard":
target["gold_idx"] = inverse[target["gold_idx"]]
elif target["kind"] == "dist":
if "candidate_indices" in target:
candidates = sorted(
zip(
(inverse[index] for index in target["candidate_indices"]),
target["probs"],
target.get("candidate_ids") or [None] * len(target["probs"]),
strict=True,
)
)
target["candidate_indices"] = [index for index, _, _ in candidates]
target["probs"] = [probability for _, probability, _ in candidates]
if target.get("candidate_ids") is not None:
target["candidate_ids"] = [
identity for _, _, identity in candidates
]
else:
target["probs"] = [target["probs"][index] for index in order]
else:
target["allowed_set"] = sorted(
inverse[index] for index in target["allowed_set"]
)
return result
def _generator(seed: int, identity: str, stream: str) -> random.Random:
digest = hashlib.sha256(f"{seed}:{stream}:{identity}".encode()).digest()
return random.Random(int.from_bytes(digest, "big")) # nosec B311 - Dataset ordering.
def _permutation_identity(record: Mapping[str, Any]) -> str:
value = {
key: record.get(key) for key in ("source", "group", "id", "state", "questions")
}
return canonical_json(value, allow_nan=False)
src/axolotl/integrations/diffusion_decision/plugin.py
+193
View File
@@ -0,0 +1,193 @@
"""Axolotl registration for typed diffusion-decision training."""
from __future__ import annotations
__ci_config_keys__ = ("diffusion_decision",)
import shutil
from collections.abc import Mapping
from pathlib import Path
from typing import Any
from axolotl.integrations.base import BasePlugin
from ._util import _value
from .args import DiffusionDecisionConfig
from .manifest import (
MANIFEST_FILENAME,
DecisionManifest,
DecisionManifestCheckpointCallback,
build_decision_manifest,
)
def _output_dir(cfg: Any) -> str | Path:
output_dir: str | Path | None = (
cfg.get("output_dir") if isinstance(cfg, Mapping) else cfg.output_dir
)
if not isinstance(output_dir, (str, Path)):
raise ValueError("diffusion_decision manifest requires output_dir")
return output_dir
class DiffusionDecisionPlugin(BasePlugin):
"""Register decision settings and defer layout handling to ``DiffusionSpec``."""
def get_input_args(self) -> str:
return "axolotl.integrations.diffusion_decision.args.DiffusionDecisionArgs"
@staticmethod
def _decision_config(cfg: Any) -> DiffusionDecisionConfig | None:
value = _value(cfg, "diffusion_decision")
if value is None or isinstance(value, DiffusionDecisionConfig):
return value
if isinstance(value, Mapping):
return DiffusionDecisionConfig.model_validate(value)
raise TypeError("diffusion_decision must be a mapping")
def register(self, cfg: dict) -> None:
if self._decision_config(cfg) is None:
return
diffusion = cfg.get("diffusion_lm")
if diffusion is None:
raise ValueError("diffusion_decision requires native diffusion_lm settings")
if _value(diffusion, "from_causal_lm", False):
raise ValueError(
"diffusion_decision does not support diffusion_lm.from_causal_lm"
)
self._validate_fixed_logical_batch(cfg)
def _require_runtime(self, cfg: Any) -> DiffusionDecisionConfig | None:
decision = self._decision_config(cfg)
if decision is None:
return None
self._validate_fixed_logical_batch(cfg)
return decision
@staticmethod
def _validate_fixed_logical_batch(cfg: Any) -> None:
batch_flattening = (
cfg.get("batch_flattening")
if isinstance(cfg, Mapping)
else getattr(cfg, "batch_flattening", False)
)
micro_batch_size = _value(cfg, "micro_batch_size")
sample_packing = bool(_value(cfg, "sample_packing", False))
logical_length = _value(cfg, "sequence_len")
if sample_packing:
if not all(
isinstance(value, int) and not isinstance(value, bool) and value > 0
for value in (logical_length, micro_batch_size)
):
raise ValueError(
"diffusion_decision sample_packing requires positive sequence_len and micro_batch_size"
)
return
if (
batch_flattening is False
and isinstance(micro_batch_size, int)
and (micro_batch_size > 1)
):
raise ValueError(
"diffusion_decision fixed logical batches require batch_flattening: true when micro_batch_size is greater than one"
)
if batch_flattening is not True:
return
physical_budget = (
logical_length * micro_batch_size
if isinstance(logical_length, int) and isinstance(micro_batch_size, int)
else None
)
if not all(
isinstance(value, int) and not isinstance(value, bool) and value > 0
for value in (logical_length, physical_budget, micro_batch_size)
):
raise ValueError(
"diffusion_decision fixed logical batches require positive logical length, physical budget, and micro batch size"
)
def load_datasets(self, cfg, preprocess: bool = False):
if self._require_runtime(cfg) is None:
return None
try:
from .datasets import load_decision_datasets
except ImportError as exc:
raise RuntimeError(
"diffusion_decision dataset integration is not installed"
) from exc
return load_decision_datasets(cfg, preprocess=preprocess)
def pre_lora_load(self, cfg, model) -> None:
decision = self._require_runtime(cfg)
if decision is None:
return
from .slot_runtime import (
merge_trainable_slot_indices,
resolve_trainable_slot_runtime,
)
runtime = resolve_trainable_slot_runtime(cfg, model, decision)
if runtime is not None:
merge_trainable_slot_indices(cfg, runtime)
def get_trainer_cls(self, cfg):
if self._require_runtime(cfg) is None:
return None
try:
from .trainer import DiffusionDecisionTrainer
except ImportError as exc:
raise RuntimeError(
"diffusion_decision trainer integration is not installed"
) from exc
return DiffusionDecisionTrainer
def post_trainer_create(self, cfg, trainer) -> None:
if self._require_runtime(cfg) is None:
return
tokenizer = getattr(trainer, "processing_class", None)
if tokenizer is None:
tokenizer = getattr(trainer, "tokenizer", None)
self._decision_manifest = build_decision_manifest(
cfg,
tokenizer=tokenizer,
model=getattr(trainer, "model", None),
)
from axolotl.utils.distributed import is_main_process
if is_main_process():
self._decision_manifest.write_to(_output_dir(cfg))
trainer.add_callback(
DecisionManifestCheckpointCallback(self._decision_manifest)
)
def post_train(self, cfg, model) -> None:
if self._require_runtime(cfg) is None:
return
from axolotl.utils.distributed import is_main_process
if not is_main_process():
return
manifest = getattr(self, "_decision_manifest", None)
if not isinstance(manifest, DecisionManifest):
manifest = build_decision_manifest(cfg, model=model)
manifest.write_to(_output_dir(cfg))
def post_lora_merge(self, cfg, adapter_path: str, output_path: str) -> None:
if self._decision_config(cfg) is None:
return
source = Path(adapter_path) / MANIFEST_FILENAME
if not source.exists():
return
DecisionManifest.from_path(source)
shutil.copy2(source, Path(output_path) / source.name)
def get_collator_cls_and_kwargs(self, cfg, is_eval: bool = False):
if self._require_runtime(cfg) is None:
return None
try:
from .training_collator import decision_collator_for_config
except ImportError as exc:
raise RuntimeError(
"diffusion_decision collator integration is not installed"
) from exc
return decision_collator_for_config(cfg, is_eval=is_eval)
src/axolotl/integrations/diffusion_decision/prepared_cache.py
+402
View File
@@ -0,0 +1,402 @@
"""Versioned local prepared-row cache for decision training."""
from __future__ import annotations
import json
import logging
import re
from dataclasses import asdict, is_dataclass
from enum import Enum
from pathlib import Path
from typing import Any, Mapping, Sequence
import tokenizers
import transformers
from huggingface_hub import hf_hub_download
from ._util import (
_value,
atomic_write_text,
canonical_sha256,
model_config_overrides,
sha256_file,
)
from .loss import decision_example_from_canvas
from .records import DecisionCanvas, OrdinalMetadata
from .slots import SlotPlan
SCHEMA_VERSION = 1
PRODUCER = "decision-prepared-rows-v2"
LOG = logging.getLogger(__name__)
def _json_default(item: Any) -> Any:
if isinstance(item, Enum):
return item.value
if is_dataclass(item) and not isinstance(item, type):
return asdict(item)
if isinstance(item, Path):
return str(item)
if hasattr(item, "to_dict"):
return item.to_dict()
return str(item)
def _sha(value: Any) -> str:
return canonical_sha256(value, default=_json_default)
def _json_value(value: Any) -> Any:
if value is None or isinstance(value, (bool, int, float, str)):
return value
if isinstance(value, Path):
return str(value)
if isinstance(value, Enum):
return _json_value(value.value)
if isinstance(value, Mapping):
return {str(key): _json_value(item) for key, item in value.items()}
if isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
return [_json_value(item) for item in value]
if hasattr(value, "model_dump"):
return _json_value(value.model_dump(mode="json"))
if hasattr(value, "to_dict"):
return _json_value(value.to_dict())
raise TypeError(f"cannot serialize decision preparation value {type(value)!r}")
def _semantic_config(cfg: Any) -> dict[str, Any]:
fields = (
"base_model",
"revision_of_model",
"model_config_type",
"seed",
"sequence_len",
"micro_batch_size",
"eval_batch_size",
"gradient_accumulation_steps",
"attn_implementation",
"sample_packing",
"eval_sample_packing",
"batch_flattening",
)
result = {
name: _json_value(value)
for name in fields
if (value := _value(cfg, name)) is not None
}
for name in ("diffusion_lm", "diffusion_decision"):
value = _value(cfg, name)
if value is not None:
result[name] = _json_value(value)
return result
def _entry_identity(
cfg: Any, selected: Sequence[tuple[str, int, Any]] | None = None
) -> list[dict[str, Any]]:
entries = []
candidates = (
selected
if selected is not None
else (
(collection, index, entry)
for collection in ("datasets", "test_datasets")
for index, entry in enumerate(_value(cfg, collection, ()) or ())
)
)
for collection, index, entry in candidates:
adapter = _value(entry, "type")
if isinstance(adapter, str) and adapter.startswith("diffusion_decision."):
entries.append(
{
"collection": collection,
"index": index,
"type": adapter,
"path": _value(entry, "path"),
"split": _value(entry, "split", "train"),
"name": _value(entry, "name"),
"revision": _value(entry, "revision"),
"data_files": _value(entry, "data_files"),
}
)
return entries
def _resolved_tokenizer_revision(cfg: Any, tokenizer: Any) -> str | None:
init_kwargs = getattr(tokenizer, "init_kwargs", None)
model_config = model_config_overrides(cfg)
for value in (
getattr(tokenizer, "_commit_hash", None),
_value(init_kwargs, "_commit_hash"),
_value(model_config, "_commit_hash"),
):
if isinstance(value, str) and value:
return value
requested_revision = _value(cfg, "revision_of_model")
if isinstance(requested_revision, str) and re.fullmatch(
r"[0-9a-fA-F]{40}", requested_revision
):
return requested_revision.lower()
return None
def _tokenizer_root(
cfg: Any, tokenizer: Any
) -> tuple[Path, dict[str, str] | None] | None:
name_or_path = getattr(tokenizer, "name_or_path", None)
if not isinstance(name_or_path, str):
return None
local_root = Path(name_or_path)
if local_root.is_dir():
return local_root, None
revision = _resolved_tokenizer_revision(cfg, tokenizer)
if revision is None:
return None
for filename in (
"tokenizer_config.json",
"tokenizer.json",
"special_tokens_map.json",
"added_tokens.json",
"chat_template.jinja",
):
try:
tokenizer_file = Path(
hf_hub_download(
name_or_path,
filename=filename,
repo_type="model",
revision=revision,
local_files_only=True,
)
)
except (OSError, ValueError):
continue
if tokenizer_file.is_file():
return tokenizer_file.parent, {
"repo_id": name_or_path,
"resolved_revision": revision,
}
return None
def identity(
cfg: Any,
tokenizer: Any,
spec: Any,
source_paths: Sequence[Path],
*,
selected_entries: Sequence[tuple[str, int, Any]] | None = None,
scope: Mapping[str, Any] | None = None,
) -> tuple[str, dict[str, Any]] | None:
backend = getattr(getattr(tokenizer, "backend_tokenizer", None), "to_str", None)
vocabulary = getattr(tokenizer, "get_vocab", None)
resolved_root = _tokenizer_root(cfg, tokenizer)
if resolved_root is None or not (callable(backend) and callable(vocabulary)):
return None
token_root, hub_identity = resolved_root
tokenizer_files = sorted(
path
for path in token_root.iterdir()
if path.is_file()
and path.name
in {
"tokenizer.json",
"tokenizer_config.json",
"special_tokens_map.json",
"added_tokens.json",
"chat_template.jinja",
}
)
if not tokenizer_files or not source_paths:
return None
try:
spec_value = asdict(spec) if is_dataclass(spec) else vars(spec) # type: ignore[arg-type]
semantic_config = _semantic_config(cfg)
except (AttributeError, TypeError, ValueError):
return None
payload: dict[str, Any] = {
"schema_version": SCHEMA_VERSION,
"producer": PRODUCER,
"producer_sha256": sha256_file(Path(__file__)),
"producer_dependencies": [
(str(path.relative_to(Path(__file__).parent)), sha256_file(path))
for path in sorted(Path(__file__).parent.rglob("*.py"))
if "__pycache__" not in path.parts
],
"config": semantic_config,
"entries": _entry_identity(cfg, selected_entries),
"scope": dict(scope or {}),
"preparation": {
"model_config": model_config_overrides(cfg),
"sample_packing": _value(cfg, "sample_packing"),
"eval_sample_packing": _value(cfg, "eval_sample_packing"),
"batch_flattening": _value(cfg, "batch_flattening"),
"eval_batch_size": _value(cfg, "eval_batch_size"),
},
"spec": spec_value,
"tokenizer": {
"files": [
(str(path.relative_to(token_root)), sha256_file(path))
for path in tokenizer_files
],
"special_ids": {
name: getattr(tokenizer, name, None)
for name in (
"pad_token_id",
"bos_token_id",
"eos_token_id",
"unk_token_id",
"mask_token_id",
)
},
"vocab_size": len(tokenizer),
"chat_template": getattr(tokenizer, "chat_template", None),
"vocab": sorted(vocabulary().items()),
"added_vocab": sorted(
getattr(tokenizer, "get_added_vocab", lambda: {})().items()
),
"backend": backend(),
"class": f"{type(tokenizer).__module__}.{type(tokenizer).__qualname__}",
"tokenizers_version": tokenizers.__version__,
"transformers_version": transformers.__version__,
},
"sources": [(str(path.resolve()), sha256_file(path)) for path in source_paths],
}
if hub_identity is not None:
payload["tokenizer"]["hub"] = hub_identity
return _sha(payload), payload
def _canvas_to_json(canvas: DecisionCanvas) -> dict[str, Any]:
data = asdict(canvas)
return data
def _canvas_from_json(data: Mapping[str, Any]) -> DecisionCanvas:
ordinal = tuple(
None
if item is None
else OrdinalMetadata(
levels=tuple(item["levels"]),
source_ids=tuple(item["source_ids"]),
candidate_ranks=tuple(item["candidate_ranks"]),
)
for item in data.get("ordinal_metadata", ())
)
return DecisionCanvas(
prompt_ids=tuple(data["prompt_ids"]),
canvas_ids=tuple(data["canvas_ids"]),
label_positions=tuple(data["label_positions"]),
allowed_ids=tuple(tuple(value) for value in data["allowed_ids"]),
question_ids=tuple(data["question_ids"]),
targets=tuple(data["targets"]),
pinned_mask=tuple(data["pinned_mask"]),
semantic_mask=tuple(data["semantic_mask"]),
slot_mask=tuple(data["slot_mask"]),
template_length=int(data["template_length"]),
prompt_slot_mask=tuple(data.get("prompt_slot_mask", ())),
ordinal_metadata=ordinal,
)
def _row_to_json(row: Mapping[str, Any]) -> dict[str, Any]:
value = {
key: item
for key, item in row.items()
if key not in {"canvas", "decision_example", "slot_plan"}
}
value["canvas"] = _canvas_to_json(row["canvas"])
example = row.get("decision_example")
if example is not None:
value["_decision_example_source_weight"] = example.source_weight
plan = row.get("slot_plan")
if plan is not None:
value["slot_plan"] = asdict(plan)
return value
def _row_from_json(data: Mapping[str, Any]) -> dict[str, Any]:
row = dict(data)
if isinstance(row.get("record"), dict) and "grouped_record_ids" in row["record"]:
row["record"] = dict(row["record"])
row["record"]["grouped_record_ids"] = tuple(row["record"]["grouped_record_ids"])
row["canvas"] = _canvas_from_json(row["canvas"])
if "slot_plan" in row:
plan = row["slot_plan"]
row["slot_plan"] = SlotPlan(
ids=tuple(plan["ids"]),
placement=plan["placement"],
pinned_mask=tuple(plan["pinned_mask"]),
update_mask=tuple(plan["update_mask"]),
loss_mask=tuple(plan["loss_mask"]),
trainable_token_ids=tuple(plan["trainable_token_ids"]),
)
source_weight = float(
row.pop("_decision_example_source_weight", row.get("source_weight", 1.0))
)
row["decision_example"] = decision_example_from_canvas(
row["canvas"], source_weight=source_weight
)
return row
def load(
root: Path,
key: str,
identity_payload: Mapping[str, Any],
) -> tuple[list[dict[str, Any]], list[dict[str, Any]], dict[str, Any]] | None:
path = root / "diffusion_decision_cache" / f"{key}.json"
try:
payload = json.loads(path.read_text())
if _sha(payload["content"]) != payload["sha256"]:
return None
metadata = payload["content"]
if _sha(metadata["identity"]) != _sha(identity_payload):
LOG.info("Decision prepared-row cache identity mismatch")
return None
rows = metadata["rows"]
train = [_row_from_json(row) for row in rows["train"]]
eval_rows = [_row_from_json(row) for row in rows["eval"]]
manifest = dict(metadata["manifest"])
if "stratified_epoch_batches" in manifest:
manifest["stratified_epoch_batches"] = tuple(
tuple(batch) for batch in manifest["stratified_epoch_batches"]
)
if _sha(rows) != metadata["rows_sha256"]:
LOG.info("Decision prepared-row cache payload validation mismatch")
return None
return train, eval_rows, manifest
except (OSError, KeyError, TypeError, ValueError, json.JSONDecodeError):
return None
def store(
root: Path,
key: str,
identity_payload: Mapping[str, Any],
cfg: Any,
train: Sequence[Mapping[str, Any]],
eval_rows: Sequence[Mapping[str, Any]],
manifest: Mapping[str, Any],
) -> None:
rows = {
"train": [_row_to_json(row) for row in train],
"eval": [_row_to_json(row) for row in eval_rows],
}
content = {
"identity": identity_payload,
"config": _semantic_config(cfg),
"manifest": manifest,
"rows_sha256": _sha(rows),
"rows": rows,
}
parent = root / "diffusion_decision_cache"
parent.mkdir(parents=True, exist_ok=True)
atomic_write_text(
parent / f"{key}.json",
json.dumps(
{"sha256": _sha(content), "content": content},
sort_keys=True,
default=_json_default,
),
)
src/axolotl/integrations/diffusion_decision/preprocessing.py
+229
View File
@@ -0,0 +1,229 @@
"""Build model-tokenized decision canvases from validated source records."""
from __future__ import annotations
import random
from collections.abc import Mapping, Sequence
from typing import Any
from .adapters.jsonl import normalize_jsonl
from .prompts import decision_prompt_ids
from .readers.base import validate_canvas
from .records import DecisionCanvas, OrdinalMetadata
from .slots import SlotPlan
from .template import (
SchemaError,
parse_decision_schema,
resolve_template,
system_text,
)
def build_decision_canvas(
tokenizer,
record: Mapping[str, Any],
prompt_ids: Sequence[int] | None = None,
*,
scaffold_ids: Sequence[int],
turn_close_id: int,
pad_id: int,
vocab_size: int,
width: int = 128,
seed: int = 0,
steps: int = 1,
noise_kind: str = "uniform",
mask_token_id: int | None = None,
slot_plan: SlotPlan | None = None,
thought_open_ids: Sequence[int] = (),
thought_close_ids: Sequence[int] = (),
include_ordinal_metadata: bool = False,
codebook: str = "vendored26",
prevalidated_record: bool = False,
) -> DecisionCanvas:
if steps < 1 or width < 2 or vocab_size < 2:
raise ValueError("steps, canvas width, and vocabulary must be positive")
if noise_kind not in {"uniform", "absorbing"}:
raise ValueError("unknown noise kind")
if noise_kind == "absorbing" and (
mask_token_id is None or not 0 <= mask_token_id < vocab_size
):
raise ValueError("absorbing noise requires a valid mask_token_id")
normalized = (
dict(record)
if prevalidated_record
else normalize_jsonl(record, codebook=codebook)
)
schema = parse_decision_schema(
{
"instructions": normalized.get("instructions"),
"questions": [
{**question, "id": key}
for key, question in normalized["questions"].items()
],
},
codebook=codebook,
)
if prompt_ids is None:
prompt_ids = decision_prompt_ids(
tokenizer, system_text(schema), normalized["state"]
)
prompt_ids = tuple(prompt_ids)
plan = _slot_plan(slot_plan)
prompt_slot_mask: tuple[bool, ...] = ()
if plan is not None and plan.placement == "prompt":
prompt_ids = (*plan.ids, *prompt_ids)
prompt_slot_mask = (True,) * len(plan.ids) + (False,) * (
len(prompt_ids) - len(plan.ids)
)
head = tuple(scaffold_ids)
if plan is not None and plan.placement == "thought":
head = (*head, *thought_open_ids, *plan.ids, *thought_close_ids)
pad_slots = () if plan is None or plan.placement != "after_turn" else plan.ids
base, slots = resolve_template(
tokenizer,
schema["questions"],
head=head,
fmt=schema["format"],
width=width - len(pad_slots),
)
ids = list(base) + [turn_close_id]
slot_positions: tuple[int, ...] = ()
if plan is not None:
if plan.placement == "thought":
start = len(scaffold_ids) + len(thought_open_ids)
slot_positions = tuple(range(start, start + len(plan.ids)))
elif plan.placement == "after_turn":
start = len(ids)
slot_positions = tuple(range(start, start + len(plan.ids)))
ids.extend(plan.ids)
if len(ids) > width:
raise SchemaError("decision canvas exceeds width")
ids.extend([pad_id] * (width - len(ids)))
positions = tuple(slot["pos"] for slot in slots)
generator = random.Random(seed) # nosec B311 - Reproduce pinned serving noise.
for position in positions:
if noise_kind == "uniform":
ids[position] = generator.randrange(vocab_size)
else:
assert mask_token_id is not None
ids[position] = mask_token_id
result = DecisionCanvas(
prompt_ids=tuple(prompt_ids),
canvas_ids=tuple(ids),
label_positions=positions,
allowed_ids=tuple(tuple(slot["label_ids"]) for slot in slots),
question_ids=tuple(question["id"] for question in schema["questions"]),
targets=tuple(
normalized["labels"][question["id"]] for question in schema["questions"]
),
pinned_mask=tuple(
(
steps > 1
and index not in positions
and not (
plan is not None
and index in slot_positions
and plan.update_mask[slot_positions.index(index)]
)
)
or (
plan is not None
and index in slot_positions
and plan.pinned_mask[slot_positions.index(index)]
)
for index in range(width)
),
semantic_mask=(True,) * width,
slot_mask=tuple(index in slot_positions for index in range(width)),
template_length=len(base),
prompt_slot_mask=prompt_slot_mask,
ordinal_metadata=(
_ordinal_metadata(
normalized, tuple(tuple(slot["label_ids"]) for slot in slots)
)
if include_ordinal_metadata
else ()
),
)
validate_canvas(result)
if any(
token >= vocab_size
for tokens in (result.prompt_ids, result.canvas_ids, *result.allowed_ids)
for token in tokens
):
raise ValueError("decision token ID exceeds the model vocabulary")
return result
def _slot_plan(slot_plan: SlotPlan | None) -> SlotPlan | None:
if slot_plan is None:
return None
widths = {
len(slot_plan.ids),
len(slot_plan.pinned_mask),
len(slot_plan.update_mask),
len(slot_plan.loss_mask),
}
if len(widths) != 1:
raise ValueError("slot plan masks must match the number of slot IDs")
if any(slot_plan.loss_mask):
raise ValueError("decision slots cannot carry label loss")
if any(
update and pinned
for update, pinned in zip(
slot_plan.update_mask, slot_plan.pinned_mask, strict=True
)
):
raise ValueError("pinned decision slots cannot be updateable")
if slot_plan.placement == "none":
if slot_plan.ids:
raise ValueError("a none slot plan cannot contain slots")
return None
if slot_plan.placement not in {"thought", "after_turn", "prompt"}:
raise ValueError(f"unknown slot placement: {slot_plan.placement!r}")
return slot_plan
def _ordinal_metadata(
normalized: Mapping[str, Any], allowed_ids: tuple[tuple[int, ...], ...]
) -> tuple[OrdinalMetadata | None, ...]:
questions = normalized["questions"]
result: list[OrdinalMetadata | None] = []
for (question_id, question), candidates in zip(
questions.items(), allowed_ids, strict=True
):
del question_id
if question.get("type") != "score":
result.append(None)
continue
levels = question.get("levels")
if not isinstance(levels, Sequence) or isinstance(levels, (str, bytes)):
raise ValueError("score ordinal metadata requires source-ordered levels")
if len(levels) != len(candidates):
raise ValueError("score levels must align with template candidates")
rendered_levels = tuple(str(level) for level in levels)
if any(not level for level in rendered_levels):
raise ValueError("score levels must be nonempty")
result.append(
OrdinalMetadata(
levels=rendered_levels,
source_ids=tuple(str(index) for index in range(len(candidates))),
candidate_ranks=tuple(range(len(candidates))),
)
)
return tuple(result)
def ordinal_metadata_for_canvas(
record: Mapping[str, Any], canvas: DecisionCanvas, *, codebook: str = "vendored26"
) -> tuple[OrdinalMetadata | None, ...]:
"""Derive score ranks from the normalized source criteria, never token IDs."""
normalized = normalize_jsonl(record, codebook=codebook)
result = _ordinal_metadata(
normalized, tuple(tuple(ids) for ids in canvas.allowed_ids)
)
if tuple(normalized["questions"]) != tuple(canvas.question_ids):
raise ValueError(
"normalized score questions must align with the prepared canvas"
)
return result
src/axolotl/integrations/diffusion_decision/prompts.py
+74
View File
@@ -0,0 +1,74 @@
"""Tokenize decision prompts with the loaded tokenizer's chat template."""
import json
from collections.abc import Mapping, Sequence
from numbers import Integral
from typing import Any
def serialize_state(state: Any) -> str:
"""Match djev's text-state serialization."""
if state is None:
raise ValueError("state is required")
return state if isinstance(state, str) else json.dumps(state, ensure_ascii=False)
def _content(text: str, structured_content: bool) -> str | list[dict[str, str]]:
if structured_content:
return [{"type": "text", "text": text}]
return text
def _input_ids(output: Any) -> tuple[int, ...]:
value = output["input_ids"] if isinstance(output, Mapping) else output
if hasattr(value, "tolist"):
value = value.tolist()
if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
raise TypeError("apply_chat_template must return token ids")
if (
value
and isinstance(value[0], Sequence)
and not isinstance(value[0], (str, bytes))
):
if len(value) != 1:
raise ValueError("apply_chat_template returned more than one prompt")
value = value[0]
if any(
isinstance(token, bool) or not isinstance(token, Integral) for token in value
):
raise TypeError("apply_chat_template must return integer token ids")
return tuple(int(token) for token in value)
def decision_prompt_ids(
tokenizer: Any,
system_text: str,
state: Any,
*,
thinking: bool = False,
structured_content: bool = False,
) -> tuple[int, ...]:
"""Return the source-compatible chat prompt ending at the model turn.
Plain text is the pinned serving contract. Callers with a tokenizer template
that requires OpenAI-style text parts must explicitly set
``structured_content=True``.
"""
state_text = serialize_state(state)
messages = [
{
"role": "system",
"content": _content(system_text, structured_content),
},
{
"role": "user",
"content": _content(state_text, structured_content),
},
]
output = tokenizer.apply_chat_template(
messages,
tokenize=True,
add_generation_prompt=True,
enable_thinking=thinking,
)
return _input_ids(output)
src/axolotl/integrations/diffusion_decision/readers/__init__.py
+10
View File
@@ -0,0 +1,10 @@
"""In-process readers for structured diffusion decisions."""
__ci_config_keys__ = ("diffusion_decision",)
from .base import DecisionRead, ReadDiagnostics
__all__ = (
"DecisionRead",
"ReadDiagnostics",
)
src/axolotl/integrations/diffusion_decision/readers/base.py
+148
View File
@@ -0,0 +1,148 @@
"""Shared typed results and validation for decision readers."""
from __future__ import annotations
from dataclasses import dataclass
from numbers import Integral
from typing import Sequence
import torch
from ..records import DecisionCanvas, OrdinalMetadata
@dataclass(frozen=True)
class ReadDiagnostics:
"""Inputs and recurrence facts retained for an optional parity fixture."""
noise_seed: int | None
noise_kind: str
update_policy: str
steps: int
forward_count: int
initial_canvas_ids: torch.Tensor
final_canvas_ids: torch.Tensor
slot_init_policy: str = "prepared_canvas_v0"
@dataclass(frozen=True)
class DecisionRead:
"""Full-vocabulary and candidate-restricted scores for one decision canvas."""
question_ids: tuple[str, ...]
label_positions: torch.Tensor
allowed_ids: torch.Tensor
candidate_mask: torch.Tensor
full_vocab_logprobs: torch.Tensor
restricted_probs: torch.Tensor
diagnostics: ReadDiagnostics | None = None
def validate_canvas(canvas: DecisionCanvas) -> int:
"""Validate the model-agnostic canvas invariants before any model execution."""
width = len(canvas.canvas_ids)
if not canvas.prompt_ids or width < 2:
raise ValueError("a decision canvas requires a prompt and token space")
if (
isinstance(canvas.template_length, bool)
or not isinstance(canvas.template_length, Integral)
or not 0 < canvas.template_length < width
):
raise ValueError("template_length must end before the turn-close token")
for tokens in (canvas.prompt_ids, canvas.canvas_ids, *canvas.allowed_ids):
if any(
isinstance(token, bool) or not isinstance(token, Integral) or token < 0
for token in tokens
):
raise ValueError("decision token IDs must be nonnegative integers")
if len(set(canvas.question_ids)) != len(canvas.question_ids) or any(
not isinstance(key, str) or not key for key in canvas.question_ids
):
raise ValueError("question IDs must be distinct nonempty strings")
fields: tuple[tuple[str, Sequence[object]], ...] = (
("pinned_mask", canvas.pinned_mask),
("semantic_mask", canvas.semantic_mask),
("slot_mask", canvas.slot_mask),
)
for name, values in fields:
if len(values) != width:
raise ValueError(f"{name} must have one value for each canvas token")
if canvas.prompt_slot_mask:
if len(canvas.prompt_slot_mask) != len(canvas.prompt_ids):
raise ValueError(
"prompt_slot_mask must have one value for each prompt token"
)
if any(not isinstance(value, bool) for value in canvas.prompt_slot_mask):
raise TypeError("prompt_slot_mask values must be bool")
question_count = len(canvas.label_positions)
if question_count == 0:
raise ValueError("a decision canvas must contain at least one label position")
if not (
len(canvas.allowed_ids) == len(canvas.question_ids) == question_count
and len(canvas.targets) in (0, question_count)
):
raise ValueError("question metadata must align with label_positions")
if canvas.ordinal_metadata and len(canvas.ordinal_metadata) != question_count:
raise ValueError("ordinal metadata must align with canvas questions")
seen_positions: set[int] = set()
for index, (position, candidates) in enumerate(
zip(canvas.label_positions, canvas.allowed_ids, strict=True)
):
if (
isinstance(position, bool)
or not isinstance(position, Integral)
or position < 0
or position >= canvas.template_length
):
raise ValueError(
f"label position {position} is outside the decision canvas"
)
if position in seen_positions:
raise ValueError("each question must have a distinct label position")
if canvas.slot_mask[position]:
raise ValueError("latent slots cannot be label positions")
if canvas.pinned_mask[position]:
raise ValueError("pinned tokens cannot be label positions")
if not canvas.semantic_mask[position]:
raise ValueError("label positions must be semantically valid")
if not candidates:
raise ValueError(f"question {index} has no allowed label token")
if len(set(candidates)) != len(candidates):
raise ValueError(f"question {index} repeats an allowed label token")
metadata = canvas.ordinal_metadata[index] if canvas.ordinal_metadata else None
if metadata is not None:
if not isinstance(metadata, OrdinalMetadata):
raise TypeError("ordinal metadata must be OrdinalMetadata or None")
metadata.check_candidates(len(candidates))
seen_positions.add(position)
return width
def restricted_probabilities(
full_vocab_logprobs: torch.Tensor, allowed_ids: Sequence[Sequence[int]]
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Apply a per-question candidate softmax without losing full-vocabulary scores."""
if full_vocab_logprobs.ndim != 2:
raise ValueError("full_vocab_logprobs must be [questions, vocabulary]")
question_count, vocab_size = full_vocab_logprobs.shape
if question_count != len(allowed_ids):
raise ValueError("allowed_ids must contain one candidate sequence per question")
maximum = max(len(ids) for ids in allowed_ids)
ids = torch.zeros(
(question_count, maximum), dtype=torch.long, device=full_vocab_logprobs.device
)
valid = torch.zeros_like(ids, dtype=torch.bool)
for row, candidates in enumerate(allowed_ids):
if any(token < 0 or token >= vocab_size for token in candidates):
raise ValueError("allowed label token ID is outside the model vocabulary")
candidate_ids = torch.as_tensor(
candidates, dtype=torch.long, device=full_vocab_logprobs.device
)
ids[row, : candidate_ids.numel()] = candidate_ids
valid[row, : candidate_ids.numel()] = True
selected = full_vocab_logprobs.gather(1, ids)
selected = selected.masked_fill(~valid, -torch.inf)
probabilities = torch.softmax(selected, dim=-1).masked_fill(~valid, 0.0)
return ids, valid, probabilities
src/axolotl/integrations/diffusion_decision/records.py
+54
View File
@@ -0,0 +1,54 @@
from dataclasses import dataclass
from numbers import Integral
from typing import Any, Sequence
@dataclass(frozen=True)
class OrdinalMetadata:
"""Source-defined score ordering for one candidate-restricted question."""
levels: Sequence[str]
source_ids: Sequence[str]
candidate_ranks: Sequence[int]
def __post_init__(self) -> None:
count = len(self.levels)
if len(self.source_ids) != count or len(self.candidate_ranks) != count:
raise ValueError("ordinal metadata must align with allowed candidates")
if (
any(not isinstance(item, str) or not item for item in self.levels)
or any(not isinstance(item, str) or not item for item in self.source_ids)
or len(set(self.source_ids)) != count
):
raise ValueError(
"ordinal levels and source IDs must be nonempty and source IDs unique"
)
if any(
isinstance(rank, bool) or not isinstance(rank, Integral)
for rank in self.candidate_ranks
):
raise ValueError("ordinal candidate ranks must be integers")
if set(self.candidate_ranks) != set(range(count)):
raise ValueError("ordinal candidate ranks must be a zero-based bijection")
def check_candidates(self, count: int) -> None:
if count < 2:
raise ValueError("ordinal metadata requires at least two candidates")
if len(self.levels) != count:
raise ValueError("ordinal metadata must align with allowed candidates")
@dataclass(frozen=True)
class DecisionCanvas:
prompt_ids: Sequence[int]
canvas_ids: Sequence[int]
label_positions: Sequence[int]
allowed_ids: Sequence[Sequence[int]]
question_ids: Sequence[str]
targets: Sequence[Any]
pinned_mask: Sequence[bool]
semantic_mask: Sequence[bool]
slot_mask: Sequence[bool]
template_length: int
prompt_slot_mask: Sequence[bool] = ()
ordinal_metadata: Sequence[OrdinalMetadata | None] = ()
src/axolotl/integrations/diffusion_decision/slot_runtime.py
+155
View File
@@ -0,0 +1,155 @@
"""Runtime resolution for learned decision-slot embedding rows."""
from __future__ import annotations
from collections.abc import Mapping, MutableMapping, Sequence
from dataclasses import dataclass
from typing import Any
from axolotl.model_support import DiffusionSpec
from ._util import (
_value,
active_control_ids,
optional_token_id,
parse_token_ids,
require_diffusion_spec,
)
from .slots import SlotInit, SlotPlan
@dataclass(frozen=True)
class TrainableSlotRuntime:
"""Validated learned/prompt slot rows and their actual input-embedding path."""
plan: SlotPlan
embedding_module_path: str
def _set_value(value: Any, name: str, item: Any) -> None:
if isinstance(value, MutableMapping):
value[name] = item
elif isinstance(value, Mapping):
raise TypeError("decision slot configuration mapping must be mutable")
else:
setattr(value, name, item)
def _embedding_path(model: Any, embedding: Any) -> str:
if not hasattr(model, "named_modules"):
raise TypeError("model must expose named_modules for trainable slot rows")
matches = [name for name, module in model.named_modules() if module is embedding]
if len(matches) != 1 or not matches[0]:
raise ValueError(
"could not determine one non-root input embedding module path for "
"peft_trainable_token_indices"
)
return matches[0]
def resolve_trainable_slot_runtime(
cfg: Any,
model: Any,
decision: Any,
*,
spec: DiffusionSpec | None = None,
) -> TrainableSlotRuntime | None:
"""Resolve learned/prompt rows from config and the model's actual embedding."""
latent = _value(decision, "latent")
if latent is None:
raise ValueError("diffusion_decision requires latent settings")
mode = _value(latent, "mode", "none")
if mode not in {"learned", "prompt"}:
return None
if not hasattr(model, "get_input_embeddings"):
raise TypeError(
"model must expose get_input_embeddings for trainable slot rows"
)
embedding = model.get_input_embeddings()
vocab_size = _value(embedding, "num_embeddings")
if (
isinstance(vocab_size, bool)
or not isinstance(vocab_size, int)
or vocab_size < 1
):
raise ValueError("model input embedding must expose a positive num_embeddings")
config = _value(model, "config")
if config is None:
raise TypeError("model must expose config for trainable slot rows")
mask_token_id = optional_token_id(_value(config, "mask_token_id"), "mask_token_id")
pad_token_id = optional_token_id(_value(config, "pad_token_id"), "pad_token_id")
control_ids = active_control_ids(config)
invalid_control_ids = sorted(
token_id for token_id in control_ids if token_id >= vocab_size
)
if invalid_control_ids:
raise ValueError(
"model config control IDs must be within the input embedding vocabulary: "
f"{invalid_control_ids}"
)
token_ids = parse_token_ids(
_value(latent, "token_ids", ()) or (), "latent.token_ids"
)
if set(token_ids) & control_ids:
raise ValueError(
"learned/prompt slot token_ids cannot use active pad, bos, eos, or mask IDs"
)
plan = SlotInit(
mode=mode,
token_ids=token_ids,
num_slots=_value(latent, "num_slots", 0),
vocab_size=vocab_size,
pad_id=pad_token_id,
spec=require_diffusion_spec(cfg) if spec is None else spec,
mask_token_id=mask_token_id,
).build()
return TrainableSlotRuntime(plan, _embedding_path(model, embedding))
def merge_trainable_slot_indices(cfg: Any, runtime: TrainableSlotRuntime) -> None:
"""Merge learned/prompt rows into PEFT's embedding-row configuration."""
configured = _value(cfg, "peft_trainable_token_indices")
additions = list(runtime.plan.trainable_token_ids)
if not additions:
return
if configured is None:
_set_value(cfg, "peft_trainable_token_indices", additions)
return
if isinstance(configured, Sequence) and not isinstance(configured, (str, bytes)):
_set_value(
cfg,
"peft_trainable_token_indices",
_merged_indices(configured, additions),
)
return
if isinstance(configured, Mapping):
existing = configured.get(runtime.embedding_module_path)
if existing is None:
raise ValueError(
"peft_trainable_token_indices mapping must include the actual input "
f"embedding path {runtime.embedding_module_path!r} before learned/prompt "
"decision slots can be merged"
)
if not isinstance(existing, Sequence) or isinstance(existing, (str, bytes)):
raise ValueError(
"peft_trainable_token_indices mapping values must be integer sequences"
)
updated = dict(configured)
updated[runtime.embedding_module_path] = _merged_indices(existing, additions)
_set_value(cfg, "peft_trainable_token_indices", updated)
return
raise ValueError(
"peft_trainable_token_indices must be an integer sequence or an embedding-path mapping"
)
def _merged_indices(existing: Sequence[Any], additions: Sequence[int]) -> list[int]:
result: list[int] = []
for value in (*existing, *additions):
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
raise ValueError(
"peft_trainable_token_indices must contain nonnegative integers"
)
if value not in result:
result.append(value)
return result
src/axolotl/integrations/diffusion_decision/slot_sampling.py
+199
View File
@@ -0,0 +1,199 @@
"""Pure count sampling and maximum-canvas slot materialization helpers."""
from __future__ import annotations
import hashlib
from dataclasses import dataclass, replace
from typing import Final
from .records import DecisionCanvas
from .slots import SlotPlan
UNIFORM_INCLUSIVE_V1: Final = "uniform_inclusive_v1"
@dataclass(frozen=True)
class DecisionDraw:
"""One pre-sharding logical decision draw with a stable sampling identity."""
index: int
epoch: int
global_draw_ordinal: int
def __post_init__(self) -> None:
_validate_nonnegative_int(self.index, "index")
_validate_nonnegative_int(self.epoch, "epoch")
_validate_nonnegative_int(self.global_draw_ordinal, "global_draw_ordinal")
def sample_slot_count(
*,
seed: int,
epoch: int,
global_draw_ordinal: int,
max_slots: int,
min_slots: int = 0,
policy: str = UNIFORM_INCLUSIVE_V1,
) -> int:
"""Return a stateless count for one logical training draw."""
_validate_nonnegative_int(seed, "seed")
_validate_nonnegative_int(epoch, "epoch")
_validate_nonnegative_int(global_draw_ordinal, "global_draw_ordinal")
_validate_nonnegative_int(min_slots, "min_slots")
_validate_nonnegative_int(max_slots, "max_slots")
if min_slots > max_slots:
raise ValueError("min_slots must not exceed max_slots")
if policy != UNIFORM_INCLUSIVE_V1:
raise ValueError(f"unsupported slot-count policy: {policy!r}")
population = max_slots - min_slots + 1
if population > 1 << 256:
raise ValueError("slot-count population exceeds the sampler domain")
limit = (1 << 256) - ((1 << 256) % population)
counter = 0
while True:
payload = f"{policy}:{seed}:{epoch}:{global_draw_ordinal}:{counter}".encode()
value = int.from_bytes(hashlib.sha256(payload).digest(), "big")
if value < limit:
return min_slots + (value % population)
counter += 1
def project_slot_plan(maximum: SlotPlan, count: int) -> SlotPlan:
"""Keep the stable prefix of a validated maximum-count plan."""
_validate_nonnegative_int(count, "count")
maximum_count = _validate_slot_plan(maximum)
if count > maximum_count:
raise ValueError("slot count exceeds the maximum slot plan")
if count == maximum_count:
return maximum
return SlotPlan(
ids=maximum.ids[:count],
placement=maximum.placement,
pinned_mask=maximum.pinned_mask[:count],
update_mask=maximum.update_mask[:count],
loss_mask=maximum.loss_mask[:count],
trainable_token_ids=maximum.trainable_token_ids[:count],
)
def project_max_canvas(
maximum_canvas: DecisionCanvas,
maximum_plan: SlotPlan,
count: int,
*,
pad_token_id: int,
padding_pinned: bool,
) -> DecisionCanvas:
"""Materialize a smaller concrete slot count from a maximum-count canvas."""
_validate_nonnegative_int(count, "count")
_validate_nonnegative_int(pad_token_id, "pad_token_id")
if not isinstance(padding_pinned, bool):
raise ValueError("padding_pinned must be a bool")
maximum_count = _validate_slot_plan(maximum_plan)
if count > maximum_count:
raise ValueError("slot count exceeds the maximum slot plan")
if count == maximum_count:
return maximum_canvas
if maximum_plan.placement == "none":
raise ValueError("a none slot plan has no count to project")
if maximum_plan.placement == "prompt":
return _project_prompt(maximum_canvas, maximum_plan, count)
positions = _canvas_slot_positions(maximum_canvas, maximum_plan)
remove = positions[count:]
if maximum_plan.placement == "thought":
if any(
position <= positions[-1] for position in maximum_canvas.label_positions
):
raise ValueError("thought slots must precede every label position")
label_positions = tuple(
position - len(remove) for position in maximum_canvas.label_positions
)
template_length = maximum_canvas.template_length - len(remove)
elif maximum_plan.placement == "after_turn":
label_positions = tuple(maximum_canvas.label_positions)
template_length = maximum_canvas.template_length
else:
raise ValueError(f"unsupported slot placement: {maximum_plan.placement!r}")
return replace(
maximum_canvas,
canvas_ids=_delete_and_extend(maximum_canvas.canvas_ids, remove, pad_token_id),
label_positions=label_positions,
pinned_mask=_delete_and_extend(
maximum_canvas.pinned_mask, remove, padding_pinned
),
semantic_mask=_delete_and_extend(maximum_canvas.semantic_mask, remove, True),
slot_mask=_delete_and_extend(maximum_canvas.slot_mask, remove, False),
template_length=template_length,
)
def _project_prompt(
canvas: DecisionCanvas, maximum: SlotPlan, count: int
) -> DecisionCanvas:
maximum_count = len(maximum.ids)
if len(canvas.prompt_slot_mask) != len(canvas.prompt_ids):
raise ValueError("prompt slots require a complete prompt_slot_mask")
if tuple(canvas.prompt_ids[:maximum_count]) != maximum.ids:
raise ValueError("prompt slot IDs must be the leading prompt IDs")
if any(not value for value in canvas.prompt_slot_mask[:maximum_count]) or any(
canvas.prompt_slot_mask[maximum_count:]
):
raise ValueError("prompt slot mask must mark only the leading slot IDs")
if any(canvas.slot_mask):
raise ValueError("prompt slots must not occupy canvas positions")
return replace(
canvas,
prompt_ids=(
*maximum.ids[:count],
*canvas.prompt_ids[maximum_count:],
),
prompt_slot_mask=(True,) * count
+ (False,) * (len(canvas.prompt_ids) - maximum_count),
)
def _canvas_slot_positions(
canvas: DecisionCanvas, maximum: SlotPlan
) -> tuple[int, ...]:
positions = tuple(index for index, value in enumerate(canvas.slot_mask) if value)
if len(positions) != len(maximum.ids):
raise ValueError("canvas slot mask does not match the maximum slot plan")
if positions != tuple(range(positions[0], positions[0] + len(positions))):
raise ValueError("canvas slot positions must be contiguous")
if tuple(canvas.canvas_ids[index] for index in positions) != maximum.ids:
raise ValueError("canvas slot IDs do not match the maximum slot plan")
return positions
def _delete_and_extend(values, remove: tuple[int, ...], padding_value):
removed = set(remove)
kept = tuple(value for index, value in enumerate(values) if index not in removed)
return (*kept, *(padding_value for _ in remove))
def _validate_slot_plan(plan: SlotPlan) -> int:
count = len(plan.ids)
if plan.placement not in {"thought", "after_turn", "prompt"}:
raise ValueError("slot sampling requires a materialized slot placement")
masks = (plan.pinned_mask, plan.update_mask, plan.loss_mask)
if any(len(mask) != count for mask in masks):
raise ValueError("slot plan masks must match slot IDs")
if any(plan.loss_mask):
raise ValueError("sampled decision slots cannot carry loss")
if any(
pinned and update
for pinned, update in zip(plan.pinned_mask, plan.update_mask, strict=True)
):
raise ValueError("pinned sampled slots cannot be updateable")
if plan.trainable_token_ids != plan.ids[: len(plan.trainable_token_ids)]:
raise ValueError("trainable slot IDs must be a stable slot prefix")
return count
def _validate_nonnegative_int(value: int, name: str) -> None:
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
raise ValueError(f"{name} must be a nonnegative integer")
src/axolotl/integrations/diffusion_decision/slots.py
+214
View File
@@ -0,0 +1,214 @@
"""Pure latent-slot initialization contracts for decision canvases."""
from __future__ import annotations
import random
from dataclasses import dataclass
from typing import Literal, Sequence
from axolotl.model_support import DiffusionNoise, DiffusionSpec
SlotMode = Literal["none", "pad", "pinned", "learned", "free", "mask", "prompt"]
SlotPlacement = Literal["none", "thought", "after_turn", "prompt"]
@dataclass(frozen=True)
class SlotPlan:
"""Initial token values and immutable mask facts for one latent-slot arm."""
ids: tuple[int, ...]
placement: SlotPlacement
pinned_mask: tuple[bool, ...]
update_mask: tuple[bool, ...]
loss_mask: tuple[bool, ...]
trainable_token_ids: tuple[int, ...]
@dataclass(frozen=True, init=False)
class SlotInit:
"""Resolve a configured slot mode without changing a model or a canvas."""
mode: SlotMode
num_slots: int
token_ids: tuple[int, ...]
vocab_size: int | None
pad_id: int | None
spec: DiffusionSpec | None
mask_token_id: int | None
def __init__(
self,
mode: SlotMode = "none",
ids: Sequence[int] = (),
*,
token_ids: Sequence[int] | None = None,
num_slots: int = 0,
vocab_size: int | None = None,
pad_id: int | None = None,
spec: DiffusionSpec | None = None,
mask_token_id: int | None = None,
) -> None:
if token_ids is not None and ids:
raise ValueError("provide either ids or token_ids, not both")
object.__setattr__(self, "mode", mode)
object.__setattr__(self, "num_slots", num_slots)
object.__setattr__(
self, "token_ids", tuple(ids if token_ids is None else token_ids)
)
object.__setattr__(self, "vocab_size", vocab_size)
object.__setattr__(self, "pad_id", pad_id)
object.__setattr__(self, "spec", spec)
object.__setattr__(self, "mask_token_id", mask_token_id)
@property
def ids(self) -> tuple[int, ...]:
"""Compatibility alias for the configured token IDs."""
return self.token_ids
def build(
self,
*,
seed: int | None = None,
generator: random.Random | None = None,
) -> SlotPlan:
"""Validate the mode and resolve its deterministic initial slot values."""
self._validate_common()
if self.mode == "none":
return SlotPlan((), "none", (), (), (), ())
assert self.vocab_size is not None
if self.mode == "pad":
assert self.pad_id is not None
self._require_no_token_ids()
return self._fixed_plan(
(self.pad_id,) * self.num_slots, placement="after_turn"
)
if self.mode == "pinned":
values = self._pinned_values()
return self._fixed_plan(values, placement="thought")
if self.mode in {"learned", "prompt"}:
self._require_exact_token_count(self.num_slots)
if len(set(self.token_ids)) != self.num_slots:
raise ValueError(f"{self.mode} slots require distinct token_ids")
return self._fixed_plan(
self.token_ids,
placement="prompt" if self.mode == "prompt" else "thought",
trainable=True,
)
if self.mode == "mask":
self._require_no_token_ids()
mask_token_id = self._absorbing_mask_token_id()
return self._fixed_plan(
(mask_token_id,) * self.num_slots, placement="thought"
)
if self.mode == "free":
self._require_no_token_ids()
assert self.spec is not None
if self.spec.noise is DiffusionNoise.ABSORBING:
values = (self._absorbing_mask_token_id(),) * self.num_slots
else:
values = self._uniform_values(seed=seed, generator=generator)
return SlotPlan(
values,
"thought",
(False,) * self.num_slots,
(True,) * self.num_slots,
(False,) * self.num_slots,
(),
)
raise ValueError(f"unknown slot mode: {self.mode!r}")
def _validate_common(self) -> None:
if self.mode not in {
"none",
"pad",
"pinned",
"learned",
"free",
"mask",
"prompt",
}:
raise ValueError(f"unknown slot mode: {self.mode!r}")
if isinstance(self.num_slots, bool) or not isinstance(self.num_slots, int):
raise ValueError("num_slots must be an integer")
if self.num_slots < 0:
raise ValueError("num_slots must be nonnegative")
if self.mode == "none":
if self.num_slots or self.token_ids:
raise ValueError("mode=none requires num_slots=0 and no token_ids")
return
if self.num_slots < 1:
raise ValueError(f"mode={self.mode} requires num_slots to be positive")
if isinstance(self.vocab_size, bool) or not isinstance(self.vocab_size, int):
raise ValueError("slot modes require an integer vocab_size")
if self.vocab_size < 1:
raise ValueError("vocab_size must be positive")
if self.mode == "pad":
self._validate_token_id(self.pad_id, "pad_id")
if not isinstance(self.spec, DiffusionSpec):
raise TypeError("slot modes require a DiffusionSpec")
for token_id in self.token_ids:
self._validate_token_id(token_id, "token_ids")
def _validate_token_id(self, token_id: int | None, name: str) -> None:
if isinstance(token_id, bool) or not isinstance(token_id, int):
raise ValueError(f"{name} must be an integer")
assert self.vocab_size is not None
if not 0 <= token_id < self.vocab_size:
raise ValueError(f"{name} must be within the model vocabulary")
def _require_no_token_ids(self) -> None:
if self.token_ids:
raise ValueError(f"mode={self.mode} does not accept token_ids")
def _require_exact_token_count(self, count: int) -> None:
if len(self.token_ids) != count:
raise ValueError(f"mode={self.mode} requires exactly {count} token_ids")
def _pinned_values(self) -> tuple[int, ...]:
if len(self.token_ids) == 1:
return self.token_ids * self.num_slots
if len(self.token_ids) != self.num_slots:
raise ValueError(
"mode=pinned requires exactly 1 token_id or exactly num_slots distinct token_ids"
)
if len(set(self.token_ids)) != self.num_slots:
raise ValueError("mode=pinned distinct token_ids must not repeat")
return self.token_ids
def _absorbing_mask_token_id(self) -> int:
assert self.spec is not None
if self.spec.noise is not DiffusionNoise.ABSORBING:
raise ValueError(f"mode={self.mode} requires absorbing diffusion noise")
self._validate_token_id(self.mask_token_id, "mask_token_id")
assert self.mask_token_id is not None
return self.mask_token_id
def _uniform_values(
self, *, seed: int | None, generator: random.Random | None
) -> tuple[int, ...]:
if seed is not None and generator is not None:
raise ValueError("free uniform slots accept either seed or generator")
if generator is None:
if isinstance(seed, bool) or not isinstance(seed, int):
raise ValueError("free uniform slots require an explicit integer seed")
generator = random.Random(seed) # nosec B311 - caller supplied deterministic seed.
assert self.vocab_size is not None
return tuple(
generator.randrange(self.vocab_size) for _ in range(self.num_slots)
)
def _fixed_plan(
self,
values: tuple[int, ...],
*,
placement: SlotPlacement,
trainable: bool = False,
) -> SlotPlan:
return SlotPlan(
values,
placement,
(True,) * self.num_slots,
(False,) * self.num_slots,
(False,) * self.num_slots,
values if trainable else (),
)
src/axolotl/integrations/diffusion_decision/template.py
+49
View File
@@ -0,0 +1,49 @@
"""Tokenizer-local access to the pinned djev schema and template functions."""
from .vendored.djev_template import (
FORMATS,
SchemaError,
answer_text,
parse_schema,
resolve_template as _resolve_template,
schedule,
system_text,
)
EXPANDED52 = tuple("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz")
def codebook_alphabet(codebook):
if codebook == "vendored26":
return None
if codebook == "expanded52":
return EXPANDED52
raise ValueError(f"unknown decision label codebook {codebook!r}")
def parse_decision_schema(value, *, codebook="vendored26"):
return parse_schema(value, alphabet=codebook_alphabet(codebook))
def resolve_template(tokenizer, qs, head=(), lead="", fmt="lines", width=128):
return _resolve_template(
qs,
list(head),
lead,
fmt,
enc=lambda text: tokenizer.encode(text, add_special_tokens=False),
canvas_length=width,
)
__all__ = [
"FORMATS",
"SchemaError",
"answer_text",
"parse_schema",
"resolve_template",
"schedule",
"system_text",
"codebook_alphabet",
"parse_decision_schema",
]
src/axolotl/integrations/diffusion_decision/trainer.py
+1512
View File
File diff suppressed because it is too large Load Diff
src/axolotl/integrations/diffusion_decision/training_collator.py
+311
View File
@@ -0,0 +1,311 @@
"""Model-agnostic batches for typed decision diffusion training."""
from __future__ import annotations
from collections.abc import Mapping, Sequence
from dataclasses import replace
from typing import Any
import torch
from axolotl.integrations.diffusion_decision._util import (
_value,
require_diffusion_spec,
)
from axolotl.integrations.diffusion_decision.collator import DecisionCanvasCollator
from axolotl.integrations.diffusion_decision.loss import (
DecisionLabelExample,
decision_example_from_canvas,
)
from axolotl.integrations.diffusion_decision.records import DecisionCanvas
from axolotl.model_support import (
DiffusionLayout,
DiffusionSpec,
)
def _decision_example(
canvas: DecisionCanvas, row: Mapping[str, Any]
) -> DecisionLabelExample:
source = row.get("decision_example")
if source is None:
source = decision_example_from_canvas(
canvas, source_weight=float(row.get("source_weight", 1.0))
)
if not isinstance(source, DecisionLabelExample):
raise TypeError("decision_example must be a DecisionLabelExample")
return replace(
source,
questions=tuple(
replace(question, position=index)
for index, question in enumerate(source.questions)
),
)
class DecisionTrainingCollator:
"""Emit layout-specific token tensors and a shared decision-loss contract."""
def __init__(
self,
tokenizer: Any | DiffusionSpec | None = None,
*,
spec: DiffusionSpec | None = None,
pad_token_id: int | None = None,
padding: bool | str = True,
max_length: int | None = None,
pad_to_multiple_of: int | None = None,
label_pad_token_id: int = -100,
return_tensors: str = "pt",
physical_payload_capacity: int | None = None,
) -> None:
if isinstance(tokenizer, DiffusionSpec):
if spec is not None:
raise TypeError("spec was provided twice")
spec = tokenizer
tokenizer = None
if not isinstance(spec, DiffusionSpec):
raise TypeError("DecisionTrainingCollator requires a DiffusionSpec")
if padding not in {True, "longest"}:
raise ValueError("DecisionTrainingCollator requires longest padding")
if label_pad_token_id != -100:
raise ValueError("DecisionTrainingCollator uses native decision masks")
if return_tensors != "pt":
raise ValueError("DecisionTrainingCollator requires return_tensors='pt'")
self.spec = spec
self.pad_token_id = int(
pad_token_id
if pad_token_id is not None
else _value(tokenizer, "pad_token_id", 0) or 0
)
self.max_length = max_length
self.pad_to_multiple_of = pad_to_multiple_of
self.physical_payload_capacity = physical_payload_capacity
self._encoder_canvas = DecisionCanvasCollator(pad_token_id=self.pad_token_id)
def __call__(self, rows: Sequence[Mapping[str, Any]]) -> dict[str, Any]:
rows = self._flatten_packed_rows(rows)
if not rows:
raise ValueError("empty decision batch")
canvases = tuple(self._canvas(row) for row in rows)
examples = tuple(
_decision_example(canvas, row)
for canvas, row in zip(canvases, rows, strict=True)
)
metadata = _decision_metadata(canvases, examples)
metadata["decision_sources"] = tuple(self._source(row) for row in rows)
metadata["decision_slot_counts"] = torch.tensor(
[
int(
row.get(
"decision_slot_count",
sum(canvas.prompt_slot_mask)
if canvas.prompt_slot_mask
else sum(canvas.slot_mask),
)
)
for row, canvas in zip(rows, canvases, strict=True)
],
dtype=torch.long,
)
metadata["decision_draws"] = tuple(row.get("decision_draw") for row in rows)
if self.spec.layout is DiffusionLayout.FULL_SEQUENCE:
return self._full_sequence(canvases, examples, metadata)
if self.spec.layout is DiffusionLayout.ENCODER_CANVAS:
result = dict(metadata)
batch = self._encoder_canvas(list(canvases))
result["diffusion_batch"] = batch
result["decision_prompt_input_pinned_mask"] = batch.encoder_validity.clone()
result["decision_logical_rows"] = metadata["decision_logical_rows"]
result["decision_label_rows"] = torch.full_like(
metadata["decision_label_rows"], -1
)
return result
raise ValueError(f"unsupported decision layout: {self.spec.layout}")
@staticmethod
def _flatten_packed_rows(
rows: Sequence[Mapping[str, Any] | Sequence[Mapping[str, Any]]],
) -> list[Mapping[str, Any]]:
"""Accept the nested logical examples emitted by MultipackBatchSampler."""
flattened: list[Mapping[str, Any]] = []
for row in rows:
if isinstance(row, Mapping):
flattened.append(row)
elif isinstance(row, Sequence) and not isinstance(row, (str, bytes)):
if not all(isinstance(item, Mapping) for item in row):
raise TypeError("packed decision batches must contain row mappings")
flattened.extend(row)
else:
raise TypeError("decision batches must contain row mappings")
return flattened
@staticmethod
def _canvas(row: Mapping[str, Any]) -> DecisionCanvas:
canvas = row.get("canvas")
if not isinstance(canvas, DecisionCanvas):
raise TypeError(
"decision dataset rows require a DecisionCanvas under 'canvas'"
)
return canvas
@staticmethod
def _source(row: Mapping[str, Any]) -> str:
source = row.get("source")
if not isinstance(source, str) or not source:
raise ValueError("decision dataset rows require a nonempty source")
return source
def _full_sequence(
self,
canvases: Sequence[DecisionCanvas],
examples: tuple[DecisionLabelExample, ...],
metadata: dict[str, Any],
) -> dict[str, Any]:
ids: list[int] = []
documents: list[int] = []
positions: list[int] = []
loss: list[bool] = []
corruptible: list[bool] = []
pinned: list[bool] = []
slots: list[bool] = []
prompt_slots: list[bool] = []
label_positions: list[list[int]] = []
offset = 0
for document, canvas in enumerate(canvases):
prompt = tuple(int(token) for token in canvas.prompt_ids)
output = tuple(int(token) for token in canvas.canvas_ids)
ids.extend((*prompt, *output))
documents.extend([document] * (len(prompt) + len(output)))
positions.extend(range(len(prompt) + len(output)))
label_set = set(int(position) for position in canvas.label_positions)
output_loss = [index in label_set for index in range(len(output))]
loss.extend([False] * len(prompt) + output_loss)
corruptible.extend([False] * len(prompt) + output_loss)
pinned.extend([True] * len(prompt) + list(canvas.pinned_mask))
slots.extend([False] * len(prompt) + list(canvas.slot_mask))
prompt_slots.extend(canvas.prompt_slot_mask or (False,) * len(prompt))
prompt_slots.extend([False] * len(output))
label_positions.append(
[offset + len(prompt) + position for position in canvas.label_positions]
)
offset += len(prompt) + len(output)
if (
self.physical_payload_capacity is not None
and len(ids) > self.physical_payload_capacity
):
raise ValueError(
"decision full-sequence microbatch exceeds native physical payload capacity"
)
input_ids = torch.tensor(ids, dtype=torch.long)[None]
metadata.update(
{
"input_ids": input_ids,
"document_ids": torch.tensor(documents, dtype=torch.long)[None],
"semantic_validity": torch.ones_like(input_ids, dtype=torch.bool),
"position_ids": torch.tensor(positions, dtype=torch.long)[None],
"canvas_loss_mask": torch.tensor(loss, dtype=torch.bool)[None],
"canvas_corruptible_mask": torch.tensor(corruptible, dtype=torch.bool)[
None
],
"canvas_input_pinned_mask": torch.tensor(pinned, dtype=torch.bool)[
None
],
"canvas_update_mask": torch.tensor(corruptible, dtype=torch.bool)[None]
& ~torch.tensor(pinned, dtype=torch.bool)[None],
"decision_slot_mask": torch.tensor(slots, dtype=torch.bool)[None],
"decision_prompt_slot_mask": torch.tensor(
prompt_slots, dtype=torch.bool
)[None],
"decision_prompt_input_pinned_mask": torch.tensor(
[True] * len(ids), dtype=torch.bool
)[None],
"decision_label_rows": _label_rows(label_positions),
"decision_label_positions": _label_positions(label_positions),
}
)
metadata["decision_examples"] = examples
return metadata
def _decision_metadata(
canvases: Sequence[DecisionCanvas], examples: tuple[DecisionLabelExample, ...]
) -> dict[str, Any]:
width = max(len(canvas.question_ids) for canvas in canvases)
batch = len(canvases)
question_mask = torch.zeros((batch, width), dtype=torch.bool)
logical_rows = torch.full((batch, width), -1, dtype=torch.long)
positions = torch.full((batch, width), -1, dtype=torch.long)
slot_mask = torch.zeros(
(batch, max(len(canvas.canvas_ids) for canvas in canvases)), dtype=torch.bool
)
prompt_slot_mask = torch.zeros(
(batch, max(len(canvas.prompt_ids) for canvas in canvases)), dtype=torch.bool
)
for row, canvas in enumerate(canvases):
count = len(canvas.question_ids)
question_mask[row, :count] = True
logical_rows[row, :count] = row
positions[row, :count] = torch.tensor(canvas.label_positions, dtype=torch.long)
slot_mask[row, : len(canvas.slot_mask)] = torch.tensor(
canvas.slot_mask, dtype=torch.bool
)
if canvas.prompt_slot_mask:
if len(canvas.prompt_slot_mask) != len(canvas.prompt_ids):
raise ValueError("prompt_slot_mask must align with prompt_ids")
prompt_slot_mask[row, : len(canvas.prompt_slot_mask)] = torch.tensor(
canvas.prompt_slot_mask, dtype=torch.bool
)
return {
"decision_examples": examples,
"decision_question_mask": question_mask,
"decision_supervision_mask": question_mask.clone(),
"decision_logical_rows": logical_rows,
"decision_label_rows": logical_rows.clone(),
"decision_label_positions": positions,
"decision_slot_mask": slot_mask,
"decision_prompt_slot_mask": prompt_slot_mask,
}
def _label_rows(label_positions: Sequence[Sequence[int]]) -> torch.Tensor:
result = torch.full(
(len(label_positions), max(map(len, label_positions))), -1, dtype=torch.long
)
for row, values in enumerate(label_positions):
result[row, : len(values)] = 0
return result
def _label_positions(label_positions: Sequence[Sequence[int]]) -> torch.Tensor:
result = torch.full(
(len(label_positions), max(map(len, label_positions))), -1, dtype=torch.long
)
for row, values in enumerate(label_positions):
result[row, : len(values)] = torch.tensor(values, dtype=torch.long)
return result
def decision_collator_for_config(cfg: Any, is_eval: bool = False):
spec = require_diffusion_spec(cfg)
tokenizer = _value(cfg, "tokenizer")
pad_token_id = int(_value(tokenizer, "pad_token_id", 0) or 0)
from axolotl.core.trainers.diffusion_lm.sampling import (
resolve_native_packing_budget,
)
packed = bool(_value(cfg, "sample_packing", False)) and (
not is_eval or _value(cfg, "eval_sample_packing", False) is not False
)
fixed_logical_batch = bool(_value(cfg, "batch_flattening", False))
budget = resolve_native_packing_budget(
cfg, packed=packed or fixed_logical_batch, is_eval=is_eval
)
return DecisionTrainingCollator, {
"spec": spec,
"pad_token_id": pad_token_id,
"physical_payload_capacity": None
if budget is None
else budget.payload_capacity,
}
src/axolotl/integrations/diffusion_decision/vendored/DJEV_LICENSE
+201
View File
@@ -0,0 +1,201 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
src/axolotl/integrations/diffusion_decision/vendored/djev_template.py
+319
View File
@@ -0,0 +1,319 @@
"""Template helpers from mmastrac/djev, Apache-2.0 (see DJEV_LICENSE).
Revision e5841cf41e9211608e698492658685c36e24e77a; source SHA-256
9348332d5c33391df50c321bc1f7721b1c57b587e333f7e80058aa6499378a24.
The resolver receives its tokenizer and width explicitly instead of globals.
"""
MAX_QUESTIONS = 64
MAX_SAMPLES = 32
SPAN_REGION = 24
SPAN_LIST_REGION = 64
class SchemaError(ValueError):
pass
def parse_schema(value, *, alphabet=None):
"""Parse a DJeV schema, with an optional local label alphabet adaptation."""
if alphabet is not None and (
not isinstance(alphabet, (list, tuple))
or len(alphabet) < 2
or any(
not isinstance(label, str)
or not label
or label != label.strip()
or "\n" in label
or "\r" in label
or len(label) != 1
for label in alphabet
)
or len(set(alphabet)) != len(alphabet)
):
raise SchemaError(
"schema: label alphabet must contain distinct one-character labels"
)
if (
not isinstance(value, dict)
or not isinstance(value.get("questions"), list)
or not value["questions"]
):
raise SchemaError("schema: needs a non-empty questions array")
if len(value["questions"]) > MAX_QUESTIONS:
raise SchemaError(f"schema: at most {MAX_QUESTIONS} questions")
qs = []
seen = set()
for q in value["questions"]:
qid = str(q.get("id", "")).strip()
if not qid or ":" in qid or "\n" in qid:
raise SchemaError(
f"question id {qid!r} must be non-empty, no ':' or newline"
)
if qid in seen:
raise SchemaError(f"duplicate question id {qid!r}")
seen.add(qid)
kind = q.get("type")
if kind in ("noul", "bool", "boolean"):
kind = "noul"
crit = q.get("criteria") or {}
choices = [("yes", crit.get("true")), ("no", crit.get("false"))]
labels = ["yes", "no"]
elif kind == "choice":
opts = q.get("options") or []
choices = [
(o["name"], o.get("description"))
if isinstance(o, dict)
else (str(o), None)
for o in opts
]
labels = (
list(alphabet[: len(choices)])
if alphabet is not None
else [chr(ord("A") + i) for i in range(len(choices))]
)
elif kind == "score":
choices = [(str(level), None) for level in (q.get("levels") or [])]
labels = (
[str(i + 1) for i in range(len(choices))]
if len(choices) <= 9
else (
list(alphabet[: len(choices)])
if alphabet is not None
else [chr(ord("A") + i) for i in range(len(choices))]
)
)
elif kind in ("span", "spans"):
default = SPAN_REGION if kind == "span" else SPAN_LIST_REGION
span_cfg = {
"max_tokens": q.get("max_tokens", default),
"max_items": q.get("max_items", 16),
}
for key, hi in (("max_tokens", 96), ("max_items", 64)):
v = span_cfg[key]
if isinstance(v, bool) or not isinstance(v, int) or not 1 <= v <= hi:
raise SchemaError(f"question {qid!r}: {key} must be 1 to {hi}")
choices, labels = [], []
else:
raise SchemaError(f"question {qid!r}: unknown type {kind!r}")
if kind not in ("span", "spans") and len(choices) < 2:
raise SchemaError(f"question {qid!r}: needs at least two alternatives")
maximum = len(alphabet) if alphabet is not None else 26
if len(choices) > maximum:
raise SchemaError(f"question {qid!r}: at most {maximum} alternatives")
deps = q.get("depends_on") or []
ask_if = q.get("ask_if") or {}
if not isinstance(deps, list) or not all(isinstance(d, str) for d in deps):
raise SchemaError(
f"question {qid!r}: depends_on must be a list of question ids"
)
if not isinstance(ask_if, dict) or not all(
isinstance(v, list) and v for v in ask_if.values()
):
raise SchemaError(
f"question {qid!r}: ask_if must map a question id to a "
"non-empty list of its answers"
)
if kind in ("span", "spans") and (deps or ask_if):
raise SchemaError(
f"question {qid!r}: a span question cannot depend on another question"
)
qs.append(
{
"id": qid,
"type": kind,
"instructions": str(q.get("instructions", "")),
"choices": choices,
"labels": labels,
"depends_on": list(dict.fromkeys(list(deps) + list(ask_if))),
"ask_if": ask_if,
"alone": bool(q.get("alone", False)),
"span": span_cfg if kind in ("span", "spans") else None,
}
)
by_id = {q["id"]: q for q in qs}
for q in qs:
for dep in q["depends_on"]:
if dep not in by_id or dep == q["id"]:
raise SchemaError(
f"question {q['id']!r}: depends on unknown question {dep!r}"
)
if by_id[dep]["span"]:
raise SchemaError(
f"question {q['id']!r}: cannot depend on the span question {dep!r}"
)
for dep, vals in q["ask_if"].items():
names = [c[0] for c in by_id[dep]["choices"]]
if any(v not in names for v in vals):
raise SchemaError(
f"question {q['id']!r}: ask_if values for {dep!r} "
f"must be among {names}"
)
schedule(qs) # refuses a cycle
samples = value.get("samples", "auto")
if samples == "auto":
policy = {
"mode": "auto",
"max": max(1, min(int(value.get("auto_max", 4)), MAX_SAMPLES)),
"threshold": float(value.get("auto_threshold", 0.1)),
}
elif isinstance(samples, int) and samples >= 1:
policy = {"mode": "fixed", "n": min(samples, MAX_SAMPLES)}
else:
raise SchemaError('schema: samples must be a positive count or "auto"')
ask = value.get("ask")
if ask is not None:
if not isinstance(ask, list) or not ask or any(a not in seen for a in ask):
raise SchemaError("schema: ask must list question ids from this schema")
for q in qs:
if q["id"] in ask and any(d not in ask for d in q["depends_on"]):
raise SchemaError(
f"schema: ask names {q['id']!r} but not everything it depends on"
)
chunk_rows = value.get("chunk_rows")
if chunk_rows is not None and (not isinstance(chunk_rows, int) or chunk_rows < 8):
raise SchemaError("schema: chunk_rows must be an integer of at least 8")
chunk_prompt = value.get("chunk_prompt", "own")
if chunk_prompt not in ("shared", "own"):
raise SchemaError('schema: chunk_prompt must be "shared" or "own"')
sequential = bool(value.get("sequential", False))
think = value.get("think", 0)
if isinstance(think, bool) or not isinstance(think, int) or not 0 <= think <= 4096:
raise SchemaError("schema: think must be a thought budget in tokens, 0 to 4096")
return {
"questions": qs,
"instructions": value.get("instructions"),
"policy": policy,
"steps": max(1, min(int(value.get("steps", 1)), 8)),
"think": think,
"ask": ask,
"chunk_rows": chunk_rows,
"chunk_prompt": chunk_prompt,
"sequential": sequential,
"format": "lines" if len([q for q in qs if not q["span"]]) <= 10 else "indexed",
}
FORMATS = {
"lines": (
"\n",
"{id}: ",
'Reply with one line per question, in this order, formatted as "id: label".',
),
"indexed": (
" ",
"{id}",
"Reply on one line with each question's id immediately followed by its "
"label, separated by single spaces.",
),
}
def system_text(schema, chunked=False):
s = (
"Answer a fixed set of questions about the state the user provides. "
"Each question lists its allowed answers; reply with exactly one label "
"per question.\n"
)
if schema.get("instructions"):
s += "\n" + str(schema["instructions"]).strip() + "\n"
for q in schema["questions"]:
s += f"\nQuestion {q['id']}: {q['instructions'].strip()}\n"
for (name, desc), label in zip(q["choices"], q["labels"], strict=False):
if q["type"] == "noul":
s += f" {label}: {str(desc).strip()}\n" if desc else f" {label}\n"
elif desc:
s += f" {label}: {name} ({str(desc).strip()})\n"
else:
s += f" {label}: {name}\n"
s += "\n" + FORMATS[schema.get("format", "lines")][2]
if chunked:
s += (
" A reply may cover only some of the questions; answer every line "
"that is present."
)
return s
def answer_text(qs, labels, fmt="lines"):
join, lead, _ = FORMATS[fmt]
return join.join(
lead.format(id=q["id"]) + q["labels"][i]
for q, i in zip(qs, labels, strict=False)
)
def resolve_template(qs, head, lead, fmt, *, enc, canvas_length):
"""Tokenize the answer template and find each question's slot. Every label
must change exactly one token, at the same position for all of a question's
labels, or this raises SchemaError. ``head`` is the token run the canvas
starts with: the empty thought block for a plain read, and empty when the
prompt already ends the thought channel. ``lead`` is the text before the
first answer: the join when earlier answers are in the prompt, so the
tokens match one joint template."""
base_labels = [0] * len(qs)
base = head + enc(lead + answer_text(qs, base_labels, fmt))
if len(base) + 1 > canvas_length:
raise SchemaError(
f"answer template is {len(base)} tokens; the canvas holds {canvas_length - 1}"
)
if len(qs) == 1 and len(base) + 1 > canvas_length:
raise SchemaError(
f"question {qs[0]['id']!r} alone needs {len(base) + 1} canvas rows"
)
slots = []
for qi, q in enumerate(qs):
pos = None
ids = [0] * len(q["labels"])
for li in range(1, len(q["labels"])):
labels = list(base_labels)
labels[qi] = li
e = head + enc(lead + answer_text(qs, labels, fmt))
if len(e) != len(base):
raise SchemaError(
f"question {q['id']!r}: label {q['labels'][li]!r} is not a "
"single token"
)
diffs = [i for i in range(len(e)) if e[i] != base[i]]
if len(diffs) != 1 or (pos is not None and diffs[0] != pos):
raise SchemaError(
f"question {q['id']!r}: labels do not share one template slot"
)
pos = diffs[0]
ids[li] = e[pos]
ids[0] = base[pos]
if len(set(ids)) != len(ids):
raise SchemaError(
f"question {q['id']!r}: two labels tokenize to the same id"
)
slots.append({"pos": pos, "label_ids": ids})
return base, slots
def schedule(qs):
"""Questions in stages: a question's stage comes after the stages of
everything it depends on. Declaration order is kept within a stage."""
ids = {q["id"] for q in qs}
pending = list(qs)
done: set = set()
levels = []
while pending:
level = [
q
for q in pending
if all(d in done or d not in ids for d in q["depends_on"])
]
if not level:
raise SchemaError(
"schema: dependency cycle among " + ", ".join(q["id"] for q in pending)
)
levels.append(level)
done |= {q["id"] for q in level}
pending = [q for q in pending if q["id"] not in done]
return levels
tests/integrations/diffusion_decision/__init__.py
View File
tests/integrations/diffusion_decision/fixtures/djev_template_e5841cf.json
+1
View File
File diff suppressed because one or more lines are too long
tests/integrations/diffusion_decision/helpers.py
+446
View File
@@ -0,0 +1,446 @@
from __future__ import annotations
import json
from pathlib import Path
from types import SimpleNamespace
from typing import Any
import torch
from torch import nn
from transformers import DiffusionGemmaConfig
from axolotl.core.trainers.diffusion_lm.backends.full_sequence import (
FullSequenceBackend,
)
from axolotl.integrations.diffusion_decision.loss import decision_example_from_canvas
from axolotl.integrations.diffusion_decision.preprocessing import build_decision_canvas
from axolotl.integrations.diffusion_decision.records import DecisionCanvas
from axolotl.integrations.diffusion_decision.slots import SlotInit, SlotPlan
from axolotl.integrations.diffusion_decision.trainer import DiffusionDecisionTrainer
from axolotl.model_support import (
DiffusionLayout,
DiffusionNoise,
DiffusionSpec,
EosHandling,
FirstPositionAlignment,
GenerationAdapter,
LogitAlignment,
MaskTokenPolicy,
ObjectiveReduction,
ReductionScope,
TimeWeighting,
)
DJEV_FIXTURE = json.loads(
(Path(__file__).parent / "fixtures" / "djev_template_e5841cf.json").read_text()
)
class CharacterTokenizer:
def encode(self, text, add_special_tokens=False):
assert not add_special_tokens
return [ord(character) for character in text]
class ChatCharacterTokenizer(CharacterTokenizer):
pad_token_id = 0
bos_token_id = 1
eos_token_id = 11
unk_token_id = 2
def apply_chat_template(
self,
messages,
*,
tokenize: bool,
add_generation_prompt: bool,
enable_thinking: bool,
) -> list[int]:
assert tokenize and add_generation_prompt and not enable_thinking
assert len(messages) == 2
return [99]
class RecordedTokenizer:
def __init__(self, encodings):
self.encodings = encodings
def encode(self, text, add_special_tokens=False):
assert not add_special_tokens
return self.encodings[text]
def make_spec(
*,
noise: DiffusionNoise = DiffusionNoise.UNIFORM,
layout: DiffusionLayout = DiffusionLayout.FULL_SEQUENCE,
self_conditioning: bool | None = None,
mask_token_policy: MaskTokenPolicy | None = None,
**fields: Any,
) -> DiffusionSpec:
values: dict[str, Any] = {
"noise": noise,
"layout": layout,
"logit_alignment": LogitAlignment.ALIGNED,
"first_position_alignment": FirstPositionAlignment.DUPLICATE_FIRST,
"self_conditioning": (
layout is DiffusionLayout.ENCODER_CANVAS
if self_conditioning is None
else self_conditioning
),
"max_canvas": 128,
"max_context": 1024,
"eos_handling": EosHandling.INDEPENDENT,
"mask_token_policy": (
(
MaskTokenPolicy.MODEL
if noise is DiffusionNoise.ABSORBING
else MaskTokenPolicy.NONE
)
if mask_token_policy is None
else mask_token_policy
),
"default_time_weighting": TimeWeighting.NONE,
"objective_reduction": ObjectiveReduction.MASKED_TOKEN_MEAN,
"generation_adapter": GenerationAdapter.FULL_SEQUENCE,
"reduction_scope": ReductionScope.MICROBATCH,
}
values.update(fields)
return DiffusionSpec(**values)
def trainer_spec(
layout: DiffusionLayout, alignment: LogitAlignment = LogitAlignment.ALIGNED
) -> DiffusionSpec:
return make_spec(
noise=DiffusionNoise.ABSORBING,
layout=layout,
logit_alignment=alignment,
max_canvas=16,
max_context=32,
objective_reduction=ObjectiveReduction.EXAMPLE_MEAN,
reduction_scope=ReductionScope.GLOBAL_WINDOW,
)
def make_record(
identifier: str | None = "record",
*,
source: str = "test",
group: str = "test",
state: Any = "state",
question: str = "q",
question_type: str = "choice",
questions: dict[str, Any] | None = None,
labels: dict[str, Any] | None = None,
**fields: Any,
) -> dict[str, Any]:
schema: dict[str, Any] = {"type": question_type}
if question_type == "choice":
schema.update(instructions="Pick.", options=["one", "two"])
record: dict[str, Any] = {
"id": identifier,
"source": source,
"group": group,
"state": state,
"questions": {question: schema} if questions is None else questions,
"labels": (
{question: {"kind": "hard", "gold_idx": 0}} if labels is None else labels
),
}
record.update(fields)
return {key: value for key, value in record.items() if value is not None}
def make_canvas(
prompt_ids=(1, 2),
canvas_ids=(3, 4, 5, 6),
label_positions=(1,),
name: str = "q",
*,
allowed=(1, 2, 3),
allowed_ids=None,
question_ids=None,
targets=None,
**fields: Any,
) -> DecisionCanvas:
width = len(canvas_ids)
values: dict[str, Any] = {
"prompt_ids": tuple(prompt_ids),
"canvas_ids": tuple(canvas_ids),
"label_positions": tuple(label_positions),
"allowed_ids": (
tuple(tuple(allowed) for _ in label_positions)
if allowed_ids is None
else allowed_ids
),
"question_ids": (
tuple(f"{name}-{index}" for index in range(len(label_positions)))
if question_ids is None
else question_ids
),
"targets": (
tuple({"kind": "hard", "gold_idx": 0} for _ in label_positions)
if targets is None
else targets
),
"pinned_mask": (False,) * width,
"semantic_mask": (True,) * width,
"slot_mask": (False,) * width,
"template_length": width,
}
values.update(fields)
return DecisionCanvas(**values)
def make_rows(canvases, *, weights=None) -> list[dict[str, Any]]:
return [
{
"canvas": canvas,
"source": f"source-{index}",
"decision_example": decision_example_from_canvas(
canvas, **({} if weights is None else {"source_weight": weights[index]})
),
}
for index, canvas in enumerate(canvases)
]
def slot_token_ids(mode: str, count: int = 2) -> tuple[int, ...]:
if mode == "pinned":
return (7,)
if mode in {"learned", "prompt"}:
return (7, 8, 10)[:count]
return ()
def free_plan_seed(mode: str, noise: DiffusionNoise) -> int | None:
return 17 if mode == "free" and noise is DiffusionNoise.UNIFORM else None
def make_slot_init(
mode,
*,
count: int = 2,
ids: tuple[int, ...] = (),
noise: DiffusionNoise = DiffusionNoise.UNIFORM,
spec: DiffusionSpec | None = None,
vocab_size: int = 256,
**fields: Any,
) -> SlotInit:
return SlotInit(
mode,
token_ids=ids,
num_slots=count,
vocab_size=vocab_size,
pad_id=0,
spec=make_spec(noise=noise) if spec is None else spec,
mask_token_id=9,
**fields,
)
def make_slot_plan(
mode,
*,
count: int = 2,
ids: tuple[int, ...] = (),
noise: DiffusionNoise = DiffusionNoise.UNIFORM,
spec: DiffusionSpec | None = None,
seed: int | None = None,
) -> SlotPlan:
return make_slot_init(mode, count=count, ids=ids, noise=noise, spec=spec).build(
seed=seed
)
def build_canvas(
slot_plan: SlotPlan | None = None,
*,
record: dict[str, Any] | None = None,
tokenizer=None,
prompt_ids=(99,),
noise: DiffusionNoise = DiffusionNoise.UNIFORM,
noise_kind: str | None = None,
steps: int = 1,
width: int = 128,
seed: int = 23,
thought_open_ids=(70, 71),
thought_close_ids=(72,),
**fields: Any,
) -> DecisionCanvas:
kind = noise.value if noise_kind is None else noise_kind
arguments: dict[str, Any] = {
"prompt_ids": prompt_ids,
"scaffold_ids": (),
"turn_close_id": 106,
"pad_id": 0,
"vocab_size": 256,
"width": width,
"seed": seed,
"steps": steps,
"noise_kind": kind,
"mask_token_id": 9 if kind == "absorbing" else None,
"slot_plan": slot_plan,
"thought_open_ids": thought_open_ids,
"thought_close_ids": thought_close_ids,
}
arguments.update(fields)
return build_decision_canvas(
CharacterTokenizer() if tokenizer is None else tokenizer,
make_record() if record is None else record,
**arguments,
)
def make_cfg(**overrides: Any) -> dict[str, Any]:
cfg: dict[str, Any] = {
"seed": 23,
"model_config": {"vocab_size": 256, "mask_token_id": 9},
"diffusion_lm": {"canvas_width": 128, "mask_token_id": 9},
"diffusion_decision": {"max_questions_per_canvas": 20},
}
for key, value in overrides.items():
if isinstance(value, dict) and isinstance(cfg.get(key), dict):
cfg[key].update(value)
else:
cfg[key] = value
return cfg
def tiny_gemma_config() -> DiffusionGemmaConfig:
return DiffusionGemmaConfig(
text_config={
"vocab_size": 32,
"hidden_size": 16,
"intermediate_size": 32,
"num_hidden_layers": 1,
"num_attention_heads": 2,
"num_key_value_heads": 1,
"head_dim": 8,
"max_position_embeddings": 32,
"layer_types": ["full_attention"],
"per_layer_config": {"0": {"head_dim": 8}},
"sliding_window": 8,
"use_bidirectional_attention": "vision",
"num_experts": 2,
"top_k_experts": 1,
"moe_intermediate_size": 16,
"pad_token_id": 0,
"eos_token_id": 1,
"bos_token_id": 2,
},
vision_config={
"model_type": "gemma4_vision",
"hidden_size": 16,
"intermediate_size": 32,
"num_hidden_layers": 1,
"num_attention_heads": 2,
"num_key_value_heads": 2,
"head_dim": 8,
"max_position_embeddings": 32,
"patch_size": 16,
"position_embedding_size": 16,
},
canvas_length=8,
boi_token_id=31,
eoi_token_id=30,
image_token_id=29,
)
class EchoModel(nn.Module):
def __init__(
self, vocab_size: int = 256, mask_token_id: int = 9, scale: float = 1.0
) -> None:
super().__init__()
self.anchor = nn.Parameter(torch.zeros(()))
self.config = SimpleNamespace(
vocab_size=vocab_size, mask_token_id=mask_token_id
)
self.scale = scale
self.calls = 0
self.seen_input_ids: torch.Tensor | None = None
def forward(self, input_ids, attention_mask, position_ids, use_cache, **kwargs):
del attention_mask, position_ids, use_cache, kwargs
self.calls += 1
self.seen_input_ids = input_ids.detach().clone()
logits = torch.nn.functional.one_hot(
input_ids, num_classes=self.config.vocab_size
).float()
return SimpleNamespace(logits=logits * self.scale + self.anchor)
class TrainerHarness(DiffusionDecisionTrainer):
def __init__(
self,
spec: DiffusionSpec,
decision: dict[str, object] | None = None,
*,
world_size: int = 1,
native_values: dict[str, object] | None = None,
mask_token_id: int = 10,
native_time: float = 0.5,
) -> None:
self._spec = spec
self.axolotl_cfg = {"diffusion_decision": decision or {}}
self.args = SimpleNamespace(world_size=world_size)
self._special_token_ids: set[int] = set()
self._native_values = native_values or {}
self._mask_token_id = mask_token_id
self._native_time = native_time
@property
def _native_spec(self) -> DiffusionSpec:
return self._spec
def _full_sequence_backend(self) -> FullSequenceBackend:
return FullSequenceBackend(
mask_token_id=self._mask_token_id, attention_backend="dense"
)
def _native_value(self, name: str, default=None):
values: dict[str, object] = {"t_eps": 0.5, "self_conditioning": None}
values.update(self._native_values)
return values.get(name, default)
def _native_unroll_settings(self) -> tuple[int, bool]:
return 1, False
@staticmethod
def _sample_native_unroll_steps(k_max: int, device: torch.device) -> int:
del k_max, device
return 1
def _sample_native_times(
self, count: int, device: torch.device, default_eps: float
) -> torch.Tensor:
del default_eps
return torch.full((count,), self._native_time, device=device)
class MultistepTrainerHarness(TrainerHarness):
_sampled_steps_value = 1
def __init__(
self,
spec: DiffusionSpec,
*,
k_max: int,
sampled_steps: int,
grad_through_steps: bool = False,
decision: dict[str, object] | None = None,
) -> None:
super().__init__(spec, decision)
self._k_max = k_max
MultistepTrainerHarness._sampled_steps_value = sampled_steps
self._grad_through_steps = grad_through_steps
def _native_unroll_settings(self) -> tuple[int, bool]:
return self._k_max, self._grad_through_steps
@staticmethod
def _sample_native_unroll_steps(k_max: int, device: torch.device) -> int:
del k_max, device
return MultistepTrainerHarness._sampled_steps_value
tests/integrations/diffusion_decision/test_adapter_dispatch.py
+59
View File
@@ -0,0 +1,59 @@
"""Public source dispatch enforces shared validation and evaluation-only data."""
import pytest
from axolotl.integrations.diffusion_decision.adapters import normalize_record
from tests.integrations.diffusion_decision.helpers import (
make_record,
)
def local_record(source="local"):
return make_record(
None, source=source, group="a", state="observed", question_type="noul"
)
def test_prefixed_and_short_names_are_equivalent():
row = local_record()
assert normalize_record(
"diffusion_decision.jsonl", row, training=True
) == normalize_record("jsonl", row, training=True)
def test_eval_corpus_cannot_enter_training_through_jsonl():
with pytest.raises(ValueError, match="reserved for evaluation"):
normalize_record("jsonl", local_record("typed_decisions"), training=True)
assert (
normalize_record("jsonl", local_record("typed_decisions"), training=False)[
"source"
]
== "typed_decisions"
)
def test_eval_corpus_adapter_rejected_before_reading_row():
with pytest.raises(ValueError, match="declared train split"):
normalize_record("typed_decisions", {}, training=True)
def test_explicit_typed_train_adapter_is_allowed():
row = {
"id": "a",
"workflow": "workflow",
"state": "{}",
"questions": '{"q":{"type":"noul"}}',
"gold": '{"q":{"probabilities":{"false":0.0,"true":1.0}}}',
}
assert (
normalize_record("typed_decisions", row, training=True, source_split="train")[
"source"
]
== "typed_decisions"
)
def test_unknown_adapter_rejected():
with pytest.raises(ValueError, match="unknown decision adapter"):
normalize_record("typo", {}, training=False)
tests/integrations/diffusion_decision/test_budgets.py
+18
View File
@@ -0,0 +1,18 @@
import pytest
from axolotl.integrations.diffusion_decision.budgets import (
decision_budget,
exceeds_budget,
)
from axolotl.model_support import DiffusionLayout
@pytest.mark.parametrize(
"layout", [DiffusionLayout.FULL_SEQUENCE, DiffusionLayout.ENCODER_CANVAS]
)
def test_fixed_canvas_cost_charges_prompt_and_full_canvas_for_both_layouts(layout):
budget = decision_budget(prompt_tokens=19, canvas_tokens=128, layout=layout)
assert budget.logical_tokens == 147
assert budget.physical_tokens == 147
assert exceeds_budget(budget, logical_limit=146, physical_limit=None) == "logical"
assert exceeds_budget(budget, logical_limit=147, physical_limit=146) == "physical"
tests/integrations/diffusion_decision/test_canvas.py
+43
View File
@@ -0,0 +1,43 @@
"""Decision canvas visibility and atomic budget contracts."""
import pytest
from axolotl.integrations.diffusion_decision.collator import DecisionCanvasCollator
from axolotl.integrations.diffusion_decision.records import DecisionCanvas
from axolotl.integrations.diffusion_decision.template import resolve_template
from tests.integrations.diffusion_decision.helpers import (
CharacterTokenizer,
)
def test_template_indexed_and_overflow():
questions = [{"id": str(index), "labels": ["a", "b"]} for index in range(64)]
_, slots = resolve_template(
CharacterTokenizer(), questions, fmt="indexed", width=256
)
assert len(slots) == 64
assert all(len(slot["label_ids"]) == 2 for slot in slots)
with pytest.raises(ValueError):
resolve_template(CharacterTokenizer(), questions, fmt="indexed", width=8)
def test_collator_preserves_prompt_and_canvas_metadata():
canvas = DecisionCanvas(
prompt_ids=[1, 2],
canvas_ids=[3, 4, 0, 0],
label_positions=[1],
allowed_ids=[[4, 5]],
question_ids=["q"],
targets=[{"kind": "hard", "gold_idx": 0}],
pinned_mask=[True, False, True, True],
semantic_mask=[True, True, True, True],
slot_mask=[False, False, False, False],
template_length=2,
)
batch = DecisionCanvasCollator()([canvas])
assert batch.canvas_loss_mask.tolist() == [[False, True, False, False]]
assert batch.canvas_input_pinned_mask.tolist() == [[True, False, True, True]]
assert batch.decoder_prefix_lengths.tolist() == [2]
assert batch.encoder_lengths.tolist() == [2]
assert batch.canvas_semantic_validity.tolist() == [[True] * 4]
tests/integrations/diffusion_decision/test_dataset_slots.py
+260
View File
@@ -0,0 +1,260 @@
"""Fixed decision slots are resolved once by dataset preparation."""
from __future__ import annotations
from copy import deepcopy
import pytest
from axolotl.integrations.diffusion_decision import datasets
from axolotl.integrations.diffusion_decision.datasets import DecisionDataset
from axolotl.integrations.diffusion_decision.preprocessing import build_decision_canvas
from axolotl.integrations.diffusion_decision.slot_sampling import DecisionDraw
from axolotl.model_support import (
DiffusionLayout,
DiffusionNoise,
DiffusionSpec,
)
from tests.integrations.diffusion_decision.helpers import (
ChatCharacterTokenizer,
make_cfg,
make_record,
make_spec,
)
def _spec(noise: DiffusionNoise = DiffusionNoise.UNIFORM, *, encoder: bool = False):
return make_spec(
noise=noise,
layout=(
DiffusionLayout.ENCODER_CANVAS if encoder else DiffusionLayout.FULL_SEQUENCE
),
self_conditioning=False,
)
def _cfg(mode: str = "none", *, count: int = 0, ids: tuple[int, ...] = ()) -> dict:
return make_cfg(
model_config={"thought_open_ids": [70], "thought_close_ids": [71]},
diffusion_decision={
"latent": {"mode": mode, "num_slots": count, "token_ids": ids}
},
)
def _prepared(cfg: dict, spec: DiffusionSpec):
rows, dropped = datasets._canvas_rows(
ChatCharacterTokenizer(), [make_record()], cfg, {}, spec
)
assert dropped == 0
assert len(rows) == 1
return rows[0]
@pytest.mark.parametrize(
("mode", "count", "ids", "noise", "placement", "expected"),
[
("pad", 2, (), DiffusionNoise.UNIFORM, "after_turn", (0, 0)),
("pinned", 2, (7,), DiffusionNoise.UNIFORM, "thought", (7, 7)),
("learned", 2, (7, 8), DiffusionNoise.UNIFORM, "thought", (7, 8)),
("prompt", 2, (7, 8), DiffusionNoise.UNIFORM, "prompt", (7, 8)),
("mask", 2, (), DiffusionNoise.ABSORBING, "thought", (9, 9)),
],
)
def test_dataset_rows_apply_fixed_slot_plan(
mode: str,
count: int,
ids: tuple[int, ...],
noise: DiffusionNoise,
placement: str,
expected: tuple[int, ...],
):
row = _prepared(_cfg(mode, count=count, ids=ids), _spec(noise))
canvas = row["canvas"]
plan = row["slot_plan"]
assert plan.placement == placement
assert plan.ids == expected
assert all(not canvas.slot_mask[position] for position in canvas.label_positions)
if placement == "prompt":
assert canvas.prompt_ids[:count] == expected
assert not any(canvas.slot_mask)
else:
positions = tuple(
index for index, selected in enumerate(canvas.slot_mask) if selected
)
assert len(positions) == count
assert tuple(canvas.canvas_ids[index] for index in positions) == expected
assert all(canvas.pinned_mask[index] for index in positions)
assert plan.trainable_token_ids == ((7, 8) if mode in {"learned", "prompt"} else ())
def test_none_mode_uses_the_legacy_canvas_call_and_row_shape():
cfg = _cfg()
record = make_record()
row = _prepared(cfg, _spec())
expected = build_decision_canvas(
ChatCharacterTokenizer(),
datasets.permute_record(deepcopy(record), seed=23),
scaffold_ids=(),
turn_close_id=11,
pad_id=0,
vocab_size=256,
width=128,
seed=23,
steps=1,
noise_kind="uniform",
mask_token_id=9,
)
assert row["canvas"] == expected
assert set(row) == {
"canvas",
"decision_example",
"source",
"length",
"record",
"max_questions",
}
def test_slot_width_overflow_is_counted_without_truncation():
cfg = _cfg("learned", count=2, ids=(7, 8))
cfg["diffusion_lm"]["canvas_width"] = 8
rows, dropped = datasets._canvas_rows(
ChatCharacterTokenizer(), [make_record()], cfg, {}, _spec()
)
assert rows == []
assert dropped == 1
@pytest.mark.parametrize(
"mode,ids", [("pinned", (7,)), ("learned", (7,)), ("prompt", (7,)), ("pad", ())]
)
def test_dynamic_slots_prepare_maximum_canvas_with_sampling_recipe(
mode: str, ids: tuple[int, ...]
):
cfg = _cfg()
cfg["diffusion_decision"]["latent"] = {
"mode": mode,
"num_slots": 1,
"token_ids": ids,
"sample_num_slots": True,
}
rows, dropped = datasets._canvas_rows(
ChatCharacterTokenizer(), [make_record()], cfg, {}, _spec()
)
assert dropped == 0
assert rows[0]["slot_sampling"]["max_slots"] == 1
assert len(rows[0]["slot_plan"].ids) == 1
def test_sampled_prompt_draw_projects_leading_prompt_slots_and_preserves_weight():
cfg = _cfg("prompt", count=3, ids=(7, 8, 10))
cfg["diffusion_decision"]["latent"]["sample_num_slots"] = True
rows, dropped = datasets._canvas_rows(
ChatCharacterTokenizer(), [make_record()], cfg, {}, _spec()
)
assert dropped == 0
maximum = rows[0]
maximum["decision_example"] = maximum["decision_example"].__class__(
questions=maximum["decision_example"].questions, source_weight=1.5
)
dataset = DecisionDataset(rows, {})
projected = [dataset[DecisionDraw(0, 0, ordinal)] for ordinal in range(64)]
by_count = {row["decision_slot_count"]: row for row in projected}
assert set(by_count) == {0, 1, 2, 3}
for count, row in by_count.items():
canvas = row["canvas"]
assert canvas.prompt_ids[:count] == (7, 8, 10)[:count]
assert canvas.prompt_slot_mask == (True,) * count + (False,) * (
len(canvas.prompt_ids) - count
)
assert not any(canvas.slot_mask)
assert row["decision_example"].source_weight == 1.5
def test_encoder_thought_slots_require_explicit_model_contract_delimiters():
cfg = _cfg("pinned", count=1, ids=(7,))
del cfg["model_config"]["thought_open_ids"]
del cfg["model_config"]["thought_close_ids"]
with pytest.raises(ValueError, match=r"model_config: \{thought_open_ids"):
datasets._canvas_rows(
ChatCharacterTokenizer(), [make_record()], cfg, {}, _spec(encoder=True)
)
@pytest.mark.parametrize("mode", ("pinned", "learned", "prompt"))
@pytest.mark.parametrize("token_id", (0, 1, 2, 11, 9, 70, 71))
def test_fixed_slot_ids_reject_active_tokenizer_model_and_delimiter_controls(
mode: str, token_id: int
):
ids = (token_id,) if mode == "pinned" else (token_id, 8)
cfg = _cfg(mode, count=2 if mode != "pinned" else 2, ids=ids)
with pytest.raises(ValueError, match="active control IDs"):
datasets._canvas_rows(
ChatCharacterTokenizer(), [make_record()], cfg, {}, _spec()
)
def test_fixed_slots_reject_tokenizer_eos_without_model_config_metadata():
cfg = _cfg("pinned", count=1, ids=(11,))
cfg["model_config"] = {"vocab_size": 256}
with pytest.raises(ValueError, match="active control IDs"):
datasets._canvas_rows(
ChatCharacterTokenizer(), [make_record()], cfg, {}, _spec()
)
def test_validated_cfg_model_config_overrides_reach_slot_plan(
min_base_cfg, monkeypatch
):
from collections import OrderedDict
from axolotl.integrations.base import PluginManager
from axolotl.integrations.diffusion_decision._util import model_config_overrides
from axolotl.integrations.diffusion_decision.plugin import (
DiffusionDecisionPlugin,
)
from axolotl.utils.config import validate_config
from axolotl.utils.dict import DictDefault
plugin_path = "axolotl.integrations.diffusion_decision.DiffusionDecisionPlugin"
monkeypatch.setattr(
PluginManager.get_instance(),
"plugins",
OrderedDict({plugin_path: DiffusionDecisionPlugin()}),
)
raw = _cfg("pinned", count=1, ids=(7,))
cfg = validate_config(
DictDefault(
min_base_cfg
| {
"plugins": [plugin_path],
"seed": raw["seed"],
"model_config": raw["model_config"],
"diffusion_lm": {"from_causal_lm": False, **raw["diffusion_lm"]},
"diffusion_decision": raw["diffusion_decision"],
}
)
)
assert cfg.model_config is None
assert model_config_overrides(cfg)["thought_open_ids"] == [70]
plan, thought_open_ids, thought_close_ids = datasets._resolve_slot_plan(
ChatCharacterTokenizer(), cfg, _spec(encoder=True), 256
)
assert plan is not None and plan.placement == "thought"
assert (thought_open_ids, thought_close_ids) == ((70,), (71,))
cfg.diffusion_decision["latent"]["token_ids"] = [70]
with pytest.raises(ValueError, match="active control IDs"):
datasets._resolve_slot_plan(
ChatCharacterTokenizer(), cfg, _spec(encoder=True), 256
)
tests/integrations/diffusion_decision/test_datasets.py
+799
View File
@@ -0,0 +1,799 @@
import json
import threading
from concurrent.futures import ThreadPoolExecutor
from types import SimpleNamespace
from typing import Any
import pytest
from transformers import PretrainedConfig
from axolotl.integrations.diffusion_decision import datasets, prepared_cache
from axolotl.integrations.diffusion_decision.args import DecisionMixtureConfig
from axolotl.integrations.diffusion_decision.loss import decision_example_from_canvas
from axolotl.integrations.diffusion_decision.records import (
DecisionCanvas,
OrdinalMetadata,
)
from axolotl.integrations.diffusion_decision.slot_sampling import (
DecisionDraw,
sample_slot_count,
)
from axolotl.integrations.diffusion_decision.slots import SlotPlan
from axolotl.model_support import DiffusionLayout
from tests.integrations.diffusion_decision.helpers import make_canvas, make_record
PREMIXED = {"premixed": True, "per_batch_stratified": False}
DEFAULT_MIXTURE = {
"weights": {"alpha": 0.8, "beta": 0.2},
"max_examples_per_source": {"alpha": 2, "beta": 2},
"per_batch_stratified": True,
}
SPEC = SimpleNamespace(
max_canvas=128,
noise=datasets.DiffusionNoise.UNIFORM,
layout=DiffusionLayout.FULL_SEQUENCE,
)
IDENTITY = {"identity": 1}
def _record(source: str, group: str, state: str, identifier: str) -> dict[str, Any]:
return make_record(
identifier, source=source, group=group, state=state, question_type="noul"
)
def _premixed(draw_id: str, state: str = "state", **fields: Any) -> dict[str, Any]:
record = _record("alpha", "same", state, "origin")
record["source_metadata"] = {
"premix": {
"draw_id": draw_id,
"origin": {"source": "alpha", "id": "origin", "group": "same"},
}
}
record.update(fields)
return record
def _entry(path: str, split: str | None = None, **fields: Any) -> dict[str, Any]:
entry = {"path": path, "type": "diffusion_decision.jsonl", **fields}
if split is not None:
entry["split"] = split
return entry
DEV_ENTRY = _entry("dev", "validation")
def _cfg(
*entries: dict[str, Any],
test_datasets: tuple[dict[str, Any], ...] = (),
mixture: dict[str, Any] | None = None,
) -> dict[str, Any]:
return {
"seed": 13,
"micro_batch_size": 2,
"datasets": list(entries),
"test_datasets": list(test_datasets),
"diffusion_decision": {
"labels": {},
"mixture": DEFAULT_MIXTURE if mixture is None else mixture,
},
}
def _patch_loader(monkeypatch, rows_by_path: dict[str, list[dict[str, Any]]]) -> None:
monkeypatch.setattr(datasets, "_rows", lambda entry: rows_by_path[entry["path"]])
monkeypatch.setattr(
datasets, "normalize_record", lambda _adapter, row, **_kwargs: dict(row)
)
monkeypatch.setattr(datasets, "load_tokenizer", lambda _cfg: range(131072))
monkeypatch.setattr(datasets, "require_diffusion_spec", lambda _cfg: SPEC)
monkeypatch.setattr(
datasets,
"_canvas_row",
lambda _tokenizer, row, _cfg, source_weight, **_kwargs: {
"canvas": SimpleNamespace(prompt_ids=(1,), canvas_ids=(2,) * 128),
"record": row,
"source": row["source"],
"source_weight": source_weight,
},
)
def _unit_canvas(**fields: Any) -> DecisionCanvas:
values: dict[str, Any] = {
"prompt_ids": (1,),
"canvas_ids": (2,),
"label_positions": (0,),
"allowed": (3,),
"question_ids": ("q",),
}
return make_canvas(**{**values, **fields})
def _cache_row(identifier: str = "x") -> dict[str, Any]:
return {"canvas": _unit_canvas(), "record": {"id": identifier}, "source": "x"}
class FakeTokenizer(SimpleNamespace):
backend_tokenizer = SimpleNamespace(to_str=lambda: "backend")
pad_token_id = bos_token_id = eos_token_id = unk_token_id = mask_token_id = 0
vocab: dict[str, int] = {"x": 0}
def __len__(self) -> int:
return getattr(self, "vocab_size", len(self.vocab))
def get_vocab(self) -> dict[str, int]:
return self.vocab
@pytest.fixture
def tokenizer_root(tmp_path):
root = tmp_path / "tokenizer"
root.mkdir()
(root / "tokenizer.json").write_text("tokenizer")
return root
@pytest.fixture
def empty_source(tmp_path):
source = tmp_path / "source.jsonl"
source.write_text("{}\n")
return source
@pytest.fixture
def two_source_cfg(monkeypatch):
_patch_loader(
monkeypatch,
{
"train": [_record("alpha", "a", "a", "a"), _record("beta", "b", "b", "b")],
"dev": [_record("gamma", "g", "g", "g")],
},
)
return _cfg(_entry("train"), test_datasets=(_entry("dev", "dev"),))
def test_loader_protects_official_test_calibration_and_ood(monkeypatch):
_patch_loader(
monkeypatch,
{
"train": [
_record("alpha", "keep", "train-keep", "a0"),
_record("alpha", "state-overlap", "shared", "a1"),
_record("beta", "family-overlap", "train-family", "b0"),
_record("beta", "keep-b", "train-b", "b1"),
],
"cal": [_record("alpha", "cal", "shared", "c0")],
"ood": [_record("beta", "family-overlap", "ood-state", "o0")],
"heldout": [_record("alpha", "heldout", "heldout-state", "h0")],
"dev": [_record("gamma", "dev", "dev-state", "d0")],
},
)
cfg = _cfg(
_entry("train", "train"),
_entry("cal", "calibration"),
_entry("ood", "ood"),
test_datasets=(DEV_ENTRY, _entry("heldout", "test")),
)
result = datasets.load_decision_datasets(cfg)
retained = [row["record"] for row in result.train_dataset]
eval_rows = [row["record"] for row in result.eval_dataset]
assert {row["id"] for row in retained} == {"a0", "b1"}
assert {row["id"] for row in eval_rows} == {"d0"}
assert result.train_dataset.manifest["protected_eval_rows"] == 3
assert result.train_dataset.manifest["dropped"]["state_overlap"] == 1
assert result.train_dataset.manifest["dropped"]["family_overlap"] == 1
def test_loader_realizes_each_capped_record_once_in_stratified_batches(monkeypatch):
records = [
_record(source, f"{tag}{index}", f"s{tag}{index}", f"{tag}{index}")
for source, tag in (("alpha", "a"), ("beta", "b"))
for index in range(3)
]
_patch_loader(
monkeypatch, {"train": records, "dev": [_record("gamma", "g", "sg", "g")]}
)
cfg = _cfg(_entry("train", "train"), test_datasets=(DEV_ENTRY,))
result = datasets.load_decision_datasets(cfg)
retained = [row["record"] for row in result.train_dataset]
manifest = result.train_dataset.manifest
assert len(retained) >= 4
assert len({row["id"] for row in retained}) == 4
assert {row["source"] for row in retained} == {"alpha", "beta"}
assert manifest["dropped"]["capped"] == 2
assert manifest["mixture_probabilities"] == {"alpha": 0.8, "beta": 0.2}
assert manifest["stratified_epoch_batches"]
assert all(
len(batch) <= manifest["stratified_micro_batch_size"]
for batch in manifest["stratified_epoch_batches"]
)
def test_loader_preserves_premixed_draw_multiplicity_without_mixture_expansion(
monkeypatch,
):
_patch_loader(
monkeypatch,
{
"train": [_premixed("draw-1"), _premixed("draw-2")],
"dev": [_record("gamma", "dev", "dev-state", "dev")],
},
)
cfg = _cfg(_entry("train", "train"), test_datasets=(DEV_ENTRY,), mixture=PREMIXED)
result = datasets.load_decision_datasets(cfg)
rows = list(result.train_dataset)
assert len(rows) == 2
assert [row["record"]["source_metadata"]["premix"]["draw_id"] for row in rows] == [
"draw-1",
"draw-2",
]
assert result.train_dataset.manifest["premixed"] is True
assert result.train_dataset.manifest["mixture_probabilities"] == {}
assert result.train_dataset.manifest["stratified_epoch_batches"] == ()
@pytest.mark.parametrize(
"kwargs",
[
{"weights": {"alpha": 1.0}},
{"temperature": 1.0},
{"max_examples_per_source": 1},
{"per_batch_stratified": True},
],
)
def test_premixed_mixture_rejects_ignored_mixture_controls(kwargs):
with pytest.raises(ValueError, match="premixed=true"):
DecisionMixtureConfig(
premixed=True, **{"per_batch_stratified": False, **kwargs}
)
def test_premixed_mixture_accepts_unweighted_unstratified_rows():
assert DecisionMixtureConfig(premixed=True, per_batch_stratified=False).premixed
@pytest.mark.parametrize(
("train", "dev", "match"),
[
pytest.param(
[
_premixed("draw-1"),
_premixed("draw-2", labels={"q": {"kind": "hard", "gold_idx": 1}}),
],
None,
"identical content",
id="changed-content-for-one-origin",
),
pytest.param(
[_premixed("draw-1"), _premixed("draw-1")],
None,
"draw_id is duplicated",
id="duplicate-draw-id",
),
pytest.param(
[_premixed("draw-1", state="shared")],
[_record("gamma", "dev", "shared", "dev")],
"overlap an evaluation split",
id="decontamination-drop",
),
],
)
def test_premixed_rejects_invalid_draws(monkeypatch, train, dev, match):
_patch_loader(monkeypatch, {"train": train, "dev": dev or []})
cfg = _cfg(
_entry("train", "train"),
test_datasets=(DEV_ENTRY,) if dev else (),
mixture=PREMIXED,
)
with pytest.raises(ValueError, match=match):
datasets.load_decision_datasets(cfg)
def test_premixed_toggle_changes_prepared_cache_config():
regular = _cfg(_entry("train"))
premixed = _cfg(_entry("train"), mixture=PREMIXED)
assert prepared_cache._semantic_config(regular) != prepared_cache._semantic_config(
premixed
)
def test_family_split_and_protected_split_names():
rows = [
_record("a", "x", "one", "a"),
_record("a", "y", "two", "b"),
_record("b", "g", "three", "c"),
]
train, dev = datasets._family_dev_split(rows, ratio=0.5)
assert train and dev
assert {datasets.family_key(row) for row in train}.isdisjoint(
{datasets.family_key(row) for row in dev}
)
assert all(
datasets._is_eval({"split": split}) for split in ("test", "calibration", "ood")
)
assert not datasets._is_eval({"split": "train"})
def test_family_split_keeps_training_data_when_one_family_would_take_it_all():
rows = [_record("a", "x", f"r{i}", "a") for i in range(4)]
train, dev = datasets._family_dev_split(rows, ratio=0.5)
assert len(train) == 4 and dev == []
rows.append(_record("a", "y", "lone", "b"))
train, dev = datasets._family_dev_split(rows, ratio=0.2)
assert len(train) + len(dev) == 5 and 0 < len(dev) < 5
def test_local_jsonl_preserves_heterogeneous_nested_records_and_glob_order(tmp_path):
(tmp_path / "01.jsonl").write_text(
'{"state":{"scenario":"one","items":[1]},"options":["a"]}\n'
)
(tmp_path / "02.jsonl").write_text(
'{"state":{"scenario":"two","items":[{"x":2}]},"options":{"yes":1}}\n'
)
rows = datasets._rows(
{"path": "json", "split": "train", "data_files": str(tmp_path / "*.jsonl")}
)
assert [row["state"]["scenario"] for row in rows] == ["one", "two"]
assert isinstance(rows[0]["options"], list)
assert isinstance(rows[1]["options"], dict)
def test_local_jsonl_reports_path_and_line_for_invalid_json(tmp_path):
path = tmp_path / "bad.jsonl"
path.write_text('{"ok":true}\n{"broken":\n')
with pytest.raises(ValueError, match=r"invalid JSONL at .*bad\.jsonl:2:"):
datasets._rows({"path": "json", "split": "train", "data_files": str(path)})
@pytest.mark.parametrize("name", ["https://example.org/data.jsonl", "local.json"])
def test_nonlocal_or_non_jsonl_inputs_use_hf_loader(monkeypatch, tmp_path, name):
if name == "local.json":
path = tmp_path / name
path.write_text('[{"value":1}]')
name = str(path)
calls = []
def load(path, **kwargs):
calls.append((path, kwargs))
return [{"value": 1}]
monkeypatch.setattr(datasets, "load_dataset", load)
assert datasets._rows({"path": "json", "data_files": name, "split": "dev"}) == [
{"value": 1}
]
assert calls == [("json", {"split": "dev", "data_files": {"dev": name}})]
def test_canvas_named_invalid_question_is_not_silently_dropped(monkeypatch):
def fail(*args, **kwargs):
raise datasets.SchemaError(
"question 'canvas_choice': labels do not share one template slot"
)
monkeypatch.setattr(datasets, "_canvas_row", fail)
with pytest.raises(datasets.SchemaError, match="canvas_choice"):
datasets._canvas_rows(range(131072), [{"source": "alpha"}], {}, {}, None)
assert datasets._is_canvas_overflow(
datasets.SchemaError("answer template is 140 tokens; the canvas holds 127")
)
def test_parallel_canvas_rows_preserves_source_order_and_overflow_accounting(
monkeypatch,
):
class InlineExecutor:
def __init__(self, **_kwargs):
pass
def __enter__(self):
return self
def __exit__(self, *_args):
return False
def map(self, fn, values, **_kwargs):
return [fn(value) for value in reversed(list(values))]
def canvas_row(_tokenizer, row, _cfg, source_weight, **_kwargs):
if row["id"] == "overflow":
raise datasets.SchemaError(
"answer template is 140 tokens; the canvas holds 127"
)
return {"id": row["id"], "source_weight": source_weight}
records = [
{"id": "first", "source": "alpha"},
{"id": "overflow", "source": "alpha"},
{"id": "last", "source": "beta"},
]
monkeypatch.setattr(datasets, "_canvas_row", canvas_row)
monkeypatch.setattr(datasets, "ThreadPoolExecutor", InlineExecutor)
monkeypatch.setattr(datasets, "_resolve_slot_plan", lambda *_args: (None, (), ()))
spec = SimpleNamespace(max_canvas=None)
serial, serial_drops = datasets._canvas_rows(
range(100), records, {"dataset_num_proc": 1}, {"alpha": 2.0}, spec
)
parallel, parallel_drops = datasets._canvas_rows(
range(100), records, {"dataset_num_proc": 2}, {"alpha": 2.0}, spec
)
assert parallel == serial
assert parallel_drops == serial_drops == 1
def test_budget_filter_precedes_source_probabilities_and_sampling(
monkeypatch, two_source_cfg
):
def canvas_row(_tokenizer, row, _cfg, source_weight, **kwargs):
return {
"canvas": SimpleNamespace(
prompt_ids=(1,) * (50 if row["source"] == "alpha" else 1),
canvas_ids=(2,) * 128,
),
"record": row,
"source": row["source"],
}
monkeypatch.setattr(datasets, "_canvas_row", canvas_row)
two_source_cfg["sequence_len"] = 150
result = datasets.load_decision_datasets(two_source_cfg)
manifest = result.train_dataset.manifest
assert manifest["budget_drops"]["train"] == {"logical": 1, "physical": 0}
assert manifest["mixture_probabilities"] == {"beta": 1.0}
assert all(row["source"] == "beta" for row in result.train_dataset)
assert sum(map(len, manifest["stratified_epoch_batches"])) == len(
result.train_dataset
)
def test_budget_filter_uses_resolved_payload_capacity(monkeypatch):
calls = []
def resolve(cfg, *, packed, batch_size=None):
calls.append(packed)
return SimpleNamespace(payload_capacity=128)
monkeypatch.setattr(datasets, "resolve_native_packing_budget", resolve)
row = {"canvas": SimpleNamespace(prompt_ids=(1,), canvas_ids=(2,) * 128)}
cfg = {
"sample_packing": False,
"batch_flattening": True,
"sequence_len": 2048,
"micro_batch_size": 1,
"diffusion_lm": {},
}
rows, drops = datasets._filter_budget_rows([row], cfg, SPEC)
assert not rows
assert drops == {"logical": 0, "physical": 1}
assert calls == [True]
def test_loader_reuses_supplied_tokenizer(monkeypatch, two_source_cfg):
def unexpected(_cfg):
raise AssertionError("must reuse the supplied tokenizer")
monkeypatch.setattr(datasets, "load_tokenizer", unexpected)
result = datasets.load_decision_datasets(two_source_cfg, tokenizer=range(131072))
assert len(result.train_dataset) > 0
def test_local_prepared_cache_reuses_typed_rows_in_exact_order(
tmp_path, tokenizer_root, monkeypatch
):
source = tmp_path / "source.jsonl"
def write_source(state: str) -> None:
source.write_text(json.dumps(_record("alpha", "a", state, "id")) + "\n")
write_source("state")
tokenizer = FakeTokenizer(
name_or_path=str(tokenizer_root), chat_template="base", vocab_size=256
)
cfg = _cfg(
_entry("json", data_files=str(source)),
mixture={
"weights": {"alpha": 1.0},
"max_examples_per_source": {"alpha": 1},
"per_batch_stratified": True,
},
)
cfg.update(
{
"dataset_prepared_path": str(tmp_path / "prepared"),
"model_config": {"vocab_size": 256, "turn_close_token_id": 2},
"diffusion_lm": {"canvas_width": 128},
}
)
calls = 0
original_rows = datasets._rows
def rows(entry):
nonlocal calls
calls += 1
return original_rows(entry)
monkeypatch.setattr(datasets, "_rows", rows)
monkeypatch.setattr(datasets, "load_tokenizer", lambda _cfg: tokenizer)
monkeypatch.setattr(datasets, "require_diffusion_spec", lambda _cfg: SPEC)
def canvas_row(_tokenizer, row, _cfg, source_weight, **_kwargs):
canvas = _unit_canvas(
canvas_ids=(2,) * 128, label_positions=(3,), allowed=(4,), template_length=1
)
return {
"canvas": canvas,
"decision_example": decision_example_from_canvas(
canvas, source_weight=source_weight
),
"record": row,
"source": row["source"],
"source_weight": source_weight,
"length": 129,
}
monkeypatch.setattr(datasets, "_canvas_row", canvas_row)
first = datasets.load_decision_datasets(cfg)
first_ids = [row["record"]["id"] for row in first.train_dataset]
first_rows = [dict(row) for row in first.train_dataset]
first_schedule = first.train_dataset.manifest["stratified_epoch_batches"]
first_calls = calls
cfg["seed"] = 14
datasets.load_decision_datasets(cfg)
assert calls > first_calls
second_cold_calls = calls
cfg["seed"] = 13
second = datasets.load_decision_datasets(cfg)
assert calls == second_cold_calls
assert [row["record"]["id"] for row in second.train_dataset] == first_ids
assert [dict(row) for row in second.train_dataset] == first_rows
assert second.train_dataset.manifest["stratified_epoch_batches"] == first_schedule
assert isinstance(second.train_dataset[0]["canvas"], DecisionCanvas)
def requires_rebuild(mutate, restore):
nonlocal calls
before = calls
mutate()
datasets.load_decision_datasets(cfg)
assert calls > before
restore()
before = calls
datasets.load_decision_datasets(cfg)
assert calls == before
requires_rebuild(lambda: write_source("changed"), lambda: write_source("state"))
requires_rebuild(
lambda: setattr(tokenizer, "chat_template", "changed"),
lambda: setattr(tokenizer, "chat_template", "base"),
)
requires_rebuild(
lambda: cfg["diffusion_lm"].update(canvas_width=129),
lambda: cfg["diffusion_lm"].update(canvas_width=128),
)
requires_rebuild(
lambda: cfg.update(micro_batch_size=3), lambda: cfg.update(micro_batch_size=2)
)
def test_prepared_cache_preserves_nested_typed_values_and_live_vocab_identity(
tokenizer_root, empty_source
):
first, second = (
prepared_cache.identity(
{},
FakeTokenizer(name_or_path=str(tokenizer_root), vocab=vocab),
SPEC,
[empty_source],
)
for vocab in ({"a": 1}, {"b": 1})
)
assert first is not None and second is not None
assert first[0] != second[0]
canvas = _unit_canvas(
slot_mask=(True,),
ordinal_metadata=(OrdinalMetadata(("low",), ("0",), (0,)),),
)
row = {
"canvas": canvas,
"decision_example": decision_example_from_canvas(canvas, source_weight=0.25),
"record": {"id": "x", "grouped_record_ids": ("x", "y")},
"source": "x",
"slot_plan": SlotPlan((9,), "thought", (True,), (False,), (False,), (9,)),
}
restored = prepared_cache._row_from_json(
json.loads(json.dumps(prepared_cache._row_to_json(row)))
)
assert restored["canvas"] == canvas
assert restored["slot_plan"] == row["slot_plan"]
assert restored["decision_example"] == row["decision_example"]
assert restored["record"] == row["record"]
def test_prepared_cache_uses_exact_cached_hub_tokenizer_snapshot(
tmp_path, empty_source, monkeypatch
):
import huggingface_hub
cache = tmp_path / "hub"
repo_id = "nvidia/Nemotron-Labs-Diffusion-8B"
revisions = ("a" * 40, "b" * 40)
for revision in revisions:
snapshot = (
cache
/ "models--nvidia--Nemotron-Labs-Diffusion-8B"
/ "snapshots"
/ revision
)
snapshot.mkdir(parents=True)
(snapshot / "tokenizer_config.json").write_text("tokenizer")
monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_CACHE", str(cache))
calls = []
hf_hub_download = prepared_cache.hf_hub_download
def local_tokenizer_file(*args, **kwargs):
calls.append((args, kwargs))
return hf_hub_download(*args, **kwargs)
monkeypatch.setattr(prepared_cache, "hf_hub_download", local_tokenizer_file)
def identity(cfg):
return prepared_cache.identity(
cfg, FakeTokenizer(name_or_path=repo_id), SPEC, [empty_source]
)
first = identity({"model_config": {"_commit_hash": revisions[0]}})
second = identity({"model_config": {"_commit_hash": revisions[1]}})
from_immutable_request = identity({"revision_of_model": revisions[0]})
from_moving_request = identity({"revision_of_model": "main"})
assert (
first is not None and second is not None and from_immutable_request is not None
)
assert first[0] != second[0]
assert from_moving_request is None
assert first[1]["tokenizer"]["hub"] == {
"repo_id": repo_id,
"resolved_revision": revisions[0],
}
assert [kwargs["local_files_only"] for _, kwargs in calls] == [True, True, True]
assert [kwargs["revision"] for _, kwargs in calls] == [*revisions, revisions[0]]
assert [kwargs["filename"] for _, kwargs in calls] == ["tokenizer_config.json"] * 3
validated_shape = identity(
{
"model_config": None,
"overrides_of_model_config": {"_commit_hash": revisions[0]},
}
)
assert validated_shape is not None
assert validated_shape[1]["tokenizer"]["hub"]["resolved_revision"] == revisions[0]
def test_prepared_cache_identity_binds_entry_roles_and_model_controls(
tokenizer_root, empty_source
):
base = {
"datasets": [_entry("json", "train", data_files=str(empty_source))],
"test_datasets": [],
"model_config": {"vocab_size": 8, "turn_close_token_id": 2},
}
role = {
**base,
"datasets": [],
"test_datasets": [{**base["datasets"][0], "split": "dev"}],
}
controls = {**base, "model_config": {"vocab_size": 8, "turn_close_token_id": 3}}
keys = [
prepared_cache.identity(
cfg, FakeTokenizer(name_or_path=str(tokenizer_root)), SPEC, [empty_source]
)[0]
for cfg in (base, role, controls)
]
assert len(set(keys)) == 3
def test_prepared_cache_repairs_corrupt_payload_and_warms(tmp_path):
row = _cache_row()
prepared_cache.store(tmp_path, "key", IDENTITY, {}, [row], [], {})
path = tmp_path / "diffusion_decision_cache" / "key.json"
payload = json.loads(path.read_text())
payload["content"]["manifest"] = {"bad": True}
path.write_text(json.dumps(payload))
assert prepared_cache.load(tmp_path, "key", IDENTITY) is None
prepared_cache.store(tmp_path, "key", IDENTITY, {}, [row], [], {})
first = prepared_cache.load(tmp_path, "key", IDENTITY)
second = prepared_cache.load(tmp_path, "key", IDENTITY)
assert first is not None and second is not None
assert first[0][0]["canvas"] == row["canvas"]
def test_prepared_cache_serializes_pretrained_model_config(tmp_path):
cfg = {"model_config": PretrainedConfig(vocab_size=8, pad_token_id=0)}
identity = {"model": cfg["model_config"]}
prepared_cache.store(tmp_path, "config", identity, cfg, [_cache_row()], [], {})
assert prepared_cache.load(tmp_path, "config", identity) is not None
def test_prepared_cache_atomic_same_key_writers_and_readers(tmp_path):
prepared_cache.store(tmp_path, "shared", IDENTITY, {}, [_cache_row("old")], [], {})
barrier = threading.Barrier(3)
def writer(identifier):
barrier.wait()
prepared_cache.store(
tmp_path, "shared", IDENTITY, {}, [_cache_row(identifier)], [], {}
)
def reader():
barrier.wait()
seen = []
for _ in range(30):
loaded = prepared_cache.load(tmp_path, "shared", IDENTITY)
assert loaded is not None
seen.append(loaded[0][0]["record"]["id"])
return seen
with ThreadPoolExecutor(max_workers=3) as pool:
futures = [
pool.submit(writer, "a"),
pool.submit(writer, "b"),
pool.submit(reader),
]
seen = futures[2].result()
futures[0].result()
futures[1].result()
assert set(seen) <= {"old", "a", "b"}
assert prepared_cache.load(tmp_path, "shared", IDENTITY) is not None
def test_decision_draw_projects_maximum_row_and_preserves_weight():
canvas = make_canvas(
prompt_ids=(1,),
canvas_ids=(7, 8, 9, 4, 0, 0),
label_positions=(3,),
allowed=(4,),
question_ids=("q",),
pinned_mask=(True, True, True, False, True, True),
slot_mask=(True, True, True, False, False, False),
template_length=4,
)
plan = SlotPlan((7, 8, 9), "thought", (True,) * 3, (False,) * 3, (False,) * 3, ())
row = {
"canvas": canvas,
"slot_plan": plan,
"slot_sampling": {
"seed": 1,
"max_slots": 3,
"pad_token_id": 0,
"padding_pinned": True,
},
"decision_example": decision_example_from_canvas(canvas, source_weight=0.25),
"source": "x",
}
dataset = datasets.DecisionDataset([row], {})
seen = {}
for ordinal in range(100):
count = sample_slot_count(
seed=1, epoch=0, global_draw_ordinal=ordinal, max_slots=3
)
seen.setdefault(count, dataset[DecisionDraw(0, 0, ordinal)])
assert set(seen) == {0, 1, 2, 3}
for count, projected in seen.items():
assert projected["decision_slot_count"] == count
assert len(projected["slot_plan"].ids) == count
assert projected["decision_example"].source_weight == 0.25
assert dataset[0] is row
tests/integrations/diffusion_decision/test_decision_grad_through.py
+135
View File
@@ -0,0 +1,135 @@
"""Native encoder-canvas graph-carry coverage for decision unrolls."""
from __future__ import annotations
from dataclasses import replace
import pytest
import torch
from peft import LoraConfig, get_peft_model
import axolotl.model_support.diffusion_gemma.modeling as gemma_modeling
from axolotl.integrations.diffusion_decision.records import DecisionCanvas
from axolotl.integrations.diffusion_decision.training_collator import (
DecisionTrainingCollator,
)
from axolotl.model_support import DiffusionLayout, LogitAlignment
from axolotl.model_support.diffusion_gemma.modeling import (
AxolotlDiffusionGemmaForBlockDiffusion,
)
from tests.integrations.diffusion_decision.helpers import (
MultistepTrainerHarness,
tiny_gemma_config,
trainer_spec,
)
def _inputs():
canvas = DecisionCanvas(
prompt_ids=(2, 3),
canvas_ids=(25, 6, 7, 8, 0, 0, 0, 0),
label_positions=(1, 3),
allowed_ids=((1, 2, 3), (4, 5, 6)),
question_ids=("q0", "q1"),
targets=({"kind": "hard", "gold_idx": 0},) * 2,
pinned_mask=(True, False, False, False, True, True, True, True),
semantic_mask=(True,) * 8,
slot_mask=(True, False, False, False, False, False, False, False),
template_length=4,
)
return DecisionTrainingCollator(
trainer_spec(DiffusionLayout.ENCODER_CANVAS, LogitAlignment.ALIGNED)
)([{"canvas": canvas, "source": "native"}])
@pytest.mark.parametrize("steps", [2, 3])
@pytest.mark.parametrize("grad_through_steps", [False, True])
@pytest.mark.parametrize("evaluate", [False, True])
def test_native_gemma_peft_grad_through_steps_carries_intermediate_graph(
monkeypatch, steps, grad_through_steps, evaluate
):
torch.manual_seed(17)
base = AxolotlDiffusionGemmaForBlockDiffusion(tiny_gemma_config())
model = get_peft_model(
base,
LoraConfig(
r=2,
lora_alpha=2,
target_modules=r"^model\.decoder\.layers\.0\.self_attn\.q_proj$",
task_type=None,
),
).train()
seen = []
states = []
original = gemma_modeling.decode_packed_canvas
def traced(*args, **kwargs):
states.append(args[1].detach().clone())
conditioning = args[5] if len(args) > 5 else None
if conditioning is not None and conditioning.requires_grad:
conditioning.retain_grad()
seen.append(conditioning)
return original(*args, **kwargs)
monkeypatch.setattr(gemma_modeling, "decode_packed_canvas", traced)
inputs = _inputs()
trainer = MultistepTrainerHarness(
trainer_spec(DiffusionLayout.ENCODER_CANVAS, LogitAlignment.ALIGNED),
k_max=steps,
sampled_steps=steps,
grad_through_steps=grad_through_steps,
decision={"latent": {"mode": "learned", "num_slots": 1}},
)
model.train(not evaluate)
with torch.set_grad_enabled(not evaluate):
loss = trainer.compute_loss(model, inputs)
carried = [value for value in seen if value is not None]
assert len(carried) == steps - 1
assert len(states) == steps
assert all(torch.equal(state, states[0]) for state in states[1:])
assert not inputs["diffusion_batch"].canvas_loss_mask[0, 0]
assert all(
value.requires_grad is (grad_through_steps and not evaluate)
for value in carried
)
if evaluate:
assert not loss.requires_grad
assert all(parameter.grad is None for parameter in model.parameters())
return
loss.backward()
assert (
any(
value.grad is not None and torch.count_nonzero(value.grad)
for value in carried
)
is grad_through_steps
)
assert any(
p.grad is not None and torch.count_nonzero(p.grad)
for n, p in model.named_parameters()
if ".lora_" in n
)
def test_grad_through_steps_rejects_non_self_conditioning_layout():
trainer = MultistepTrainerHarness(
trainer_spec(DiffusionLayout.FULL_SEQUENCE, LogitAlignment.ALIGNED),
k_max=2,
sampled_steps=2,
grad_through_steps=True,
)
with pytest.raises(NotImplementedError, match="encoder-canvas self-conditioning"):
trainer.compute_loss(torch.nn.Linear(1, 1), _inputs())
def test_grad_through_steps_rejects_encoder_without_self_conditioning():
spec = replace(
trainer_spec(DiffusionLayout.ENCODER_CANVAS, LogitAlignment.ALIGNED),
self_conditioning=False,
)
trainer = MultistepTrainerHarness(
spec, k_max=2, sampled_steps=2, grad_through_steps=True
)
with pytest.raises(NotImplementedError, match="encoder-canvas self-conditioning"):
trainer.compute_loss(torch.nn.Linear(1, 1), _inputs())
tests/integrations/diffusion_decision/test_decision_unroll.py
+39
View File
@@ -0,0 +1,39 @@
"""CPU contract for the non-self-conditioning decision recurrence."""
from __future__ import annotations
import pytest
import torch
from axolotl.core.trainers.diffusion_lm.unroll import run_unroll
@pytest.mark.parametrize("steps", [1, 2, 3])
def test_decision_recurrence_updates_only_free_slots_between_reads(steps):
state = torch.tensor([[9, 4, 7]])
update_mask = torch.tensor([[True, False, False]])
seen = []
def forward(current, *_args):
seen.append(current.clone())
logits = torch.nn.functional.one_hot((current + 1) % 16, 16).float()
return type("Output", (), {"logits": logits})()
_, _, final = run_unroll(
state=state,
update_mask=update_mask,
steps=steps,
grad_through_steps=False,
supports_self_conditioning=False,
k1_conditioning_mask=None,
recurrent_conditioning_mask=None,
forward_step=forward,
logits_from_outputs=lambda output: output.logits,
update_state=lambda current, logits, mask: torch.where(
mask, logits.argmax(-1), current
),
pilot_for_single_step=False,
)
assert len(seen) == steps
assert torch.equal(final[0, 1:], state[0, 1:])
assert final.tolist() == [[9 + steps - 1, 4, 7]]
tests/integrations/diffusion_decision/test_djev_parity.py
+70
View File
@@ -0,0 +1,70 @@
"""Exact upstream fixtures generated with pinned Gemma and Dream tokenizers."""
import json
import re
import pytest
from axolotl.integrations.diffusion_decision import template
from tests.integrations.diffusion_decision.helpers import (
DJEV_FIXTURE,
RecordedTokenizer,
)
def json_value(value):
return json.loads(json.dumps(value))
@pytest.mark.parametrize("family", ["gemma", "dream"])
@pytest.mark.parametrize("index", range(22))
def test_schema_system_answer_and_resolver_match_pinned_djev(family, index):
reference = DJEV_FIXTURE["tokenizers"][family]
case = reference["cases"][index]
schema = template.parse_schema(case["schema"])
assert json_value(schema) == case["parsed"]
assert template.system_text(schema, chunked=case["chunked"]) == case["system_text"]
assert (
template.answer_text(
schema["questions"], case["answer_indices"], schema["format"]
)
== case["answer_text"]
)
args = {
"tokenizer": RecordedTokenizer(reference["encodings"]),
"qs": schema["questions"],
"head": case["head"],
"lead": case["lead"],
"fmt": schema["format"],
"width": case["width"],
}
if "template_error" in case:
with pytest.raises(
template.SchemaError, match="^" + re.escape(case["template_error"]) + "$"
):
template.resolve_template(**args)
else:
assert json_value(template.resolve_template(**args)) == case["template"]
@pytest.mark.parametrize("case", DJEV_FIXTURE["invalid_schemas"])
def test_schema_rejection_matches_pinned_djev(case):
with pytest.raises(
template.SchemaError, match="^" + re.escape(case["error"]) + "$"
):
template.parse_schema(case["schema"])
def test_fixture_includes_successful_native_tokenization_and_boundary_cases():
assert DJEV_FIXTURE["source_revision"] == "e5841cf41e9211608e698492658685c36e24e77a"
for reference in DJEV_FIXTURE["tokenizers"].values():
assert sum("template" in case for case in reference["cases"]) >= 20
assert {case["parsed"]["format"] for case in reference["cases"]} == {
"lines",
"indexed",
}
assert any(
len(case["parsed"]["questions"]) == 64 for case in reference["cases"]
)
assert reference["tokenizer_files_sha256"]
tests/integrations/diffusion_decision/test_fixed_slot_runtime.py
+247
View File
@@ -0,0 +1,247 @@
"""Fixed decision slots retain their prepared masks through training and reads."""
from __future__ import annotations
from dataclasses import replace
from typing import Any, cast
import pytest
import torch
from axolotl.integrations.diffusion_decision.loss import decision_label_loss
from axolotl.integrations.diffusion_decision.trainer import DiffusionDecisionTrainer
from axolotl.integrations.diffusion_decision.training_collator import (
DecisionTrainingCollator,
)
from axolotl.model_support import (
DiffusionNoise,
DiffusionSpec,
LogitAlignment,
ObjectiveReduction,
ReductionScope,
)
from tests.integrations.diffusion_decision.helpers import (
EchoModel,
TrainerHarness,
build_canvas,
make_record,
make_slot_plan,
make_spec,
slot_token_ids,
)
def _spec(noise: DiffusionNoise) -> DiffusionSpec:
return make_spec(
noise=noise,
objective_reduction=ObjectiveReduction.EXAMPLE_MEAN,
reduction_scope=ReductionScope.GLOBAL_WINDOW,
)
def _prepared_canvas(mode: str, noise: DiffusionNoise):
spec = _spec(noise)
plan = make_slot_plan(mode, ids=slot_token_ids(mode), spec=spec)
canvas = build_canvas(
plan,
record=make_record("fixed-slot", question="choice"),
prompt_ids=(90,),
noise=noise,
thought_open_ids=(70,),
thought_close_ids=(71,),
)
return spec, plan, canvas
def _slot_config(mode: str) -> dict[str, object]:
return {
"mode": mode,
"num_slots": 2,
"token_ids": list(slot_token_ids(mode)),
}
@pytest.mark.parametrize(
("mode", "noise"),
[
("pad", DiffusionNoise.UNIFORM),
("pinned", DiffusionNoise.UNIFORM),
("learned", DiffusionNoise.UNIFORM),
("prompt", DiffusionNoise.UNIFORM),
("mask", DiffusionNoise.ABSORBING),
],
)
def test_prepared_fixed_slots_stay_pinned_through_trainer(mode, noise):
spec, plan, canvas = _prepared_canvas(mode, noise)
fixed_canvas = torch.tensor(canvas.slot_mask, dtype=torch.bool)
label_mask = torch.zeros(len(canvas.canvas_ids), dtype=torch.bool)
label_mask[torch.tensor(canvas.label_positions)] = True
assert not torch.any(fixed_canvas & label_mask)
batch = DecisionTrainingCollator(spec)([{"canvas": canvas, "source": "test"}])
assert batch["decision_label_positions"][0, 0].item() == (
len(canvas.prompt_ids) + canvas.label_positions[0]
)
fixed_positions = (
torch.arange(len(plan.ids))
if mode == "prompt"
else torch.where(fixed_canvas)[0] + len(canvas.prompt_ids)
)
batch["canvas_corruptible_mask"][0, fixed_positions] = True
batch["canvas_update_mask"][0, fixed_positions] = True
trainer = TrainerHarness(
spec, {"latent": _slot_config(mode)}, mask_token_id=9, native_time=1.0
)
trainer._validate_decision_config(trainer._decision_config(), k_max=1, spec=spec)
train_model = EchoModel()
trainer._full_sequence_logits(train_model, batch, spec, trainer._decision_config())
assert train_model.seen_input_ids is not None
expected = torch.tensor(
(
plan.ids
if mode == "prompt"
else torch.tensor(canvas.canvas_ids)[fixed_canvas].tolist()
),
dtype=torch.long,
)
torch.testing.assert_close(train_model.seen_input_ids[0, fixed_positions], expected)
@pytest.mark.parametrize("steps", [1, 2, 3])
def test_tiny_native_dream_mask_slots_are_pinned_and_train_through_lora(
monkeypatch: pytest.MonkeyPatch, steps: int
):
from peft import LoraConfig, get_peft_model
from transformers import AutoConfig
from axolotl.model_support.dream import _model_class
from tests.native_source_fixtures import native_source_fixture_path
source = native_source_fixture_path("dream")
if source is None:
pytest.skip("native Dream source fixture unavailable")
config = AutoConfig.from_pretrained(
source, trust_remote_code=True, local_files_only=True
)
for name, value in {
"vocab_size": 256,
"hidden_size": 64,
"intermediate_size": 128,
"num_hidden_layers": 1,
"num_attention_heads": 4,
"num_key_value_heads": 2,
"head_dim": 16,
"max_position_embeddings": 128,
"max_window_layers": 1,
"bos_token_id": 1,
"eos_token_id": 1,
"pad_token_id": 0,
"mask_token_id": 9,
"use_cache": False,
}.items():
setattr(config, name, value)
config._name_or_path = str(source)
model = (
cast(Any, _model_class()).from_config(config, trust_remote_code=True).train()
)
model = get_peft_model(
model, LoraConfig(r=2, lora_alpha=4, target_modules=["q_proj"])
)
spec, plan, canvas = _prepared_canvas("mask", DiffusionNoise.ABSORBING)
spec = replace(spec, logit_alignment=LogitAlignment.SHIFTED)
batch = DecisionTrainingCollator(spec)([{"canvas": canvas, "source": "test"}])
slots = torch.where(torch.tensor(canvas.slot_mask))[0] + len(canvas.prompt_ids)
assert batch["canvas_input_pinned_mask"][0, slots].all()
assert not batch["canvas_loss_mask"][0, slots].any()
assert not batch["canvas_corruptible_mask"][0, slots].any()
assert not batch["canvas_update_mask"][0, slots].any()
seen: list[torch.Tensor] = []
raw_logits: list[torch.Tensor] = []
selected_inputs: list[torch.Tensor] = []
selected_outputs: list[torch.Tensor] = []
def capture_input(_module, args, kwargs):
input_ids = kwargs["input_ids"] if "input_ids" in kwargs else args[0]
seen.append(input_ids.detach().clone())
def capture_output(_module, _args, _kwargs, output):
raw_logits.append(output.logits.detach().clone())
original_select = DiffusionDecisionTrainer._select_question_logits
def capture_selected(logits, inputs, *, coordinates=None):
selected_inputs.append(logits.detach().clone())
selected, supervision = original_select(logits, inputs, coordinates=coordinates)
selected_outputs.append(selected.detach().clone())
return selected, supervision
monkeypatch.setattr(
DiffusionDecisionTrainer,
"_select_question_logits",
staticmethod(capture_selected),
)
base_model = model.get_base_model()
input_hook = base_model.register_forward_pre_hook(capture_input, with_kwargs=True)
output_hook = base_model.register_forward_hook(capture_output, with_kwargs=True)
class KStepHarness(TrainerHarness):
@staticmethod
def _native_unroll_settings():
return steps, False
@staticmethod
def _sample_native_unroll_steps(k_max, device):
del k_max, device
return steps
trainer = KStepHarness(
spec, {"latent": _slot_config("mask")}, mask_token_id=9, native_time=1.0
)
try:
loss = cast(torch.Tensor, trainer.compute_loss(model, batch))
finally:
input_hook.remove()
output_hook.remove()
assert torch.isfinite(loss)
assert len(seen) == len(raw_logits) == steps
assert selected_inputs and selected_outputs
assert all(
torch.equal(value[0, slots], torch.full_like(slots, 9)) for value in seen
)
labels = torch.tensor(canvas.label_positions) + len(canvas.prompt_ids)
assert all(
torch.equal(value[0, labels], torch.full_like(labels, 9)) for value in seen
)
pinned = torch.tensor(canvas.pinned_mask, dtype=torch.bool)
scaffold = torch.where(pinned)[0] + len(canvas.prompt_ids)
expected_scaffold = torch.tensor(canvas.canvas_ids)[pinned]
assert all(torch.equal(value[0, scaffold], expected_scaffold) for value in seen)
torch.testing.assert_close(seen[-1][0, slots], torch.full_like(slots, 9))
expected_aligned = torch.cat([raw_logits[-1][:, :1], raw_logits[-1][:, :-1]], dim=1)
torch.testing.assert_close(selected_inputs[-1], expected_aligned)
rows = batch["decision_label_rows"].long()
positions = batch["decision_label_positions"].long()
expected_selected = expected_aligned[rows.clamp_min(0), positions.clamp_min(0)]
torch.testing.assert_close(selected_outputs[-1], expected_selected)
decision = trainer._decision_config()
expected_loss = decision_label_loss(
selected_outputs[-1],
batch["decision_examples"],
batch["decision_supervision_mask"],
label_softmax=decision.labels.label_softmax,
brier_weight=decision.labels.brier_weight,
).loss
torch.testing.assert_close(loss.detach(), expected_loss.detach())
loss.backward()
lora_grads = [
parameter.grad
for name, parameter in model.named_parameters()
if "lora_" in name
]
assert lora_grads and all(
gradient is not None and torch.isfinite(gradient).all()
for gradient in lora_grads
)
assert any(gradient.abs().sum() > 0 for gradient in lora_grads)
tests/integrations/diffusion_decision/test_grouping.py
+112
View File
@@ -0,0 +1,112 @@
import pytest
from axolotl.integrations.diffusion_decision.grouping import group_records
def _record(
identifier, *, source="source", group="group", state='{"a": 1}', instructions=""
):
return {
"id": identifier,
"source": source,
"group": group,
"state": state,
"instructions": instructions,
"questions": {"q": {"type": "choice", "options": ["a", "b"]}},
"labels": {"q": {"kind": "hard", "gold_idx": 0 if identifier == "one" else 1}},
"source_metadata": {"origin": identifier},
}
def test_groups_identical_context_and_stably_renames_question_collisions():
grouped = group_records([_record("one"), _record("two")], max_questions=4)
assert len(grouped) == 1
record = grouped[0]
assert list(record["questions"]) == ["q", "q__2"]
assert record["labels"]["q"] == {"kind": "hard", "gold_idx": 0}
assert record["labels"]["q__2"] == {"kind": "hard", "gold_idx": 1}
assert record["grouped_record_ids"] == ("one", "two")
assert record["grouped_question_metadata"]["q__2"] == {
"record_id": "two",
"original_question_id": "q",
"source_metadata": {"origin": "two"},
}
def test_keeps_mismatched_state_source_group_or_global_instructions_separate():
records = [
_record("one"),
_record("two", state='{"a":1}'),
_record("three", source="other"),
_record("four", group="other"),
_record("five", instructions="global"),
{**_record("six"), "family": "other"},
]
grouped = group_records(records, max_questions=4)
assert len(grouped) == len(records)
assert [record["id"] for record in grouped] == [record["id"] for record in records]
assert [record["source_metadata"] for record in grouped] == [
record["source_metadata"] for record in records
]
def test_rejects_empty_question_mappings():
record = _record("one")
record["questions"] = {}
record["labels"] = {}
try:
group_records([record], max_questions=2)
except ValueError as error:
assert "at least one question" in str(error)
else:
raise AssertionError("empty questions were accepted")
def test_splits_at_question_bound_without_truncating_targets_or_metadata():
first = _record("one")
first["questions"] = {f"q{index}": {"type": "noul"} for index in range(3)}
first["labels"] = {
f"q{index}": {"kind": "dist", "probs": [index, 3 - index]} for index in range(3)
}
second = _record("two")
second["questions"] = {"x": {"type": "score"}, "y": {"type": "choice"}}
second["labels"] = {
"x": {"kind": "set", "allowed_set": [0, 1]},
"y": {"kind": "hard", "gold_idx": 1},
}
grouped = group_records([first, second], max_questions=2)
assert [len(record["questions"]) for record in grouped] == [2, 2, 1]
targets = [target for record in grouped for target in record["labels"].values()]
assert targets == [
{"kind": "dist", "probs": [0, 3]},
{"kind": "dist", "probs": [1, 2]},
{"kind": "dist", "probs": [2, 1]},
{"kind": "set", "allowed_set": [0, 1]},
{"kind": "hard", "gold_idx": 1},
]
def test_rejects_duplicate_record_ids_within_a_group():
records = [_record("one"), _record("one"), _record("two")]
with pytest.raises(ValueError, match="duplicated: one"):
group_records(records, max_questions=1)
def test_records_without_ids_get_positional_ids_and_keep_their_metadata():
records = [_record("one"), _record("two")]
for record in records:
del record["id"]
grouped = group_records(records, max_questions=1)
assert [record["id"] for record in grouped] == [
"group#record0#group0",
"group#record0#group1",
]
assert grouped[1]["grouped_record_ids"] == ("group#record1",)
assert grouped[1]["grouped_source_metadata"] == {"group#record1": {"origin": "two"}}
tests/integrations/diffusion_decision/test_hygiene.py
+62
View File
@@ -0,0 +1,62 @@
"""Leakage checks use state identity, not labels or source-specific row IDs."""
import pytest
from axolotl.integrations.diffusion_decision.hygiene import (
assert_split_isolation,
decontaminate,
state_fingerprint,
)
def row(state, family="family-a", source="source-a"):
return {"state": state, "family": family, "source": source, "labels": {}}
def test_canonical_state_identity():
assert state_fingerprint('{ "b": 2, "a": 1 }') == state_fingerprint(
{"a": 1, "b": 2}
)
assert state_fingerprint("cafe\u0301 \n value") == state_fingerprint("café value")
assert state_fingerprint({"value": [1, 2]}) != state_fingerprint({"value": [2, 1]})
def test_cross_source_decontamination_and_family_exclusion():
evaluation = [row({"x": 1})]
candidates = [
row('{"x":1}', "other", "source-b"),
row({"x": 2}),
row({"x": 3}, "family-b"),
]
kept, counts = decontaminate(iter(candidates), iter(evaluation))
assert kept == [candidates[2]]
assert counts == {"input": 3, "state_overlap": 1, "family_overlap": 1, "kept": 1}
assert candidates[0]["labels"] == {}
def test_official_split_state_filter_can_preserve_shared_families():
candidates = [row("different state")]
kept, counts = decontaminate(
candidates, [row("evaluation")], exclude_families=False
)
assert kept == candidates
assert counts["family_overlap"] == 0
@pytest.mark.parametrize("second", [row("same", "other", "source-b"), row("different")])
def test_split_guard_rejects_state_or_family_overlap(second):
with pytest.raises(ValueError, match="overlaps splits"):
assert_split_isolation({"train": [row("same")], "test": [second]})
def test_family_names_are_source_scoped_and_repetitions_within_split_are_valid():
assert_split_isolation(
{"train": [row("a"), row("b")], "test": [row("c", source="source-b")]}
)
def test_missing_state_or_family_is_not_silently_accepted():
with pytest.raises(ValueError, match="state"):
state_fingerprint(None)
with pytest.raises(ValueError, match="family"):
decontaminate([{"source": "x", "state": "y"}], [])
tests/integrations/diffusion_decision/test_jsonl_adapter.py
+56
View File
@@ -0,0 +1,56 @@
"""Normalized records reject targets incompatible with their question schema."""
from copy import deepcopy
import pytest
from axolotl.integrations.diffusion_decision.adapters.jsonl import normalize_jsonl
def record(target):
return {
"source": "local",
"group": "family-1",
"state": {"observation": 1},
"questions": {"q": {"type": "choice", "options": ["act", "wait"]}},
"labels": {"q": target},
}
@pytest.mark.parametrize(
"target",
[
{"kind": "hard", "gold_idx": 1},
{"kind": "dist", "probs": [0.2, 0.8]},
{"kind": "set", "allowed_set": [0, 1]},
],
)
def test_all_label_kinds_roundtrip_without_mutating_input(target):
source = record(target)
original = deepcopy(source)
result = normalize_jsonl(source)
assert result == original
result["questions"]["q"]["options"].append("other")
assert source == original
@pytest.mark.parametrize(
"target",
[
{"kind": "hard", "gold_idx": -1},
{"kind": "hard", "gold_idx": True},
{"kind": "set", "allowed_set": []},
{"kind": "set", "allowed_set": [0, 0]},
{"kind": "dist", "probs": [0.8, 0.8]},
],
)
def test_invalid_targets_are_rejected(target):
with pytest.raises(ValueError):
normalize_jsonl(record(target))
def test_question_target_keys_must_match():
source = record({"kind": "hard", "gold_idx": 0})
source["labels"] = {}
with pytest.raises(ValueError, match="exactly one"):
normalize_jsonl(source)
tests/integrations/diffusion_decision/test_jsonl_normalization.py
+48
View File
@@ -0,0 +1,48 @@
from __future__ import annotations
from copy import deepcopy
import pytest
from axolotl.integrations.diffusion_decision import adapters
from axolotl.integrations.diffusion_decision.adapters import (
jsonl as jsonl_adapter,
normalize_record,
)
def _row():
return {
"id": "row",
"source": "local",
"group": "group",
"state": {"value": 1},
"questions": {
"q": {"type": "choice", "instructions": "pick", "options": ["a", "b"]}
},
"labels": {"q": {"kind": "hard", "gold_idx": 1}},
}
def test_jsonl_dispatch_validates_once_and_does_not_mutate(monkeypatch):
row = _row()
original = deepcopy(row)
calls = 0
real = jsonl_adapter.normalize_jsonl
def counted(*args, **kwargs):
nonlocal calls
calls += 1
return real(*args, **kwargs)
monkeypatch.setattr(adapters, "normalize_jsonl", counted)
assert normalize_record("jsonl", row, training=False) == original
assert row == original
assert calls == 1
def test_jsonl_dispatch_preserves_malformed_error():
row = _row()
row["labels"] = {}
with pytest.raises(ValueError, match="exactly one target"):
normalize_record("jsonl", row, training=False)
tests/integrations/diffusion_decision/test_loss.py
+742
View File
@@ -0,0 +1,742 @@
import math
import pytest
import torch
import torch.nn.functional as F
from axolotl.integrations.diffusion_decision.loss import (
DecisionLabelExample,
DecisionLabelQuestion,
DistributionLabel,
HardLabel,
SetLabel,
decision_label_loss,
decision_label_loss_from_hidden,
label_target_from_mapping,
)
ONE = torch.tensor([[True]])
MIXED_MASK = torch.tensor([[True, False, False], [False, True, True]])
FULL = {"label_softmax": "full"}
RESULT_FIELDS = (
"loss",
"restricted_loss",
"full_vocab_loss",
"effective_full_vocab_loss",
"full_vocab_dft_hard_weight_sum",
"full_vocab_dft_hard_count",
"brier_loss",
)
def _dense_linear_token_loss(hidden, head, targets):
return F.cross_entropy(
F.linear(hidden, head.weight, head.bias).float(), targets, reduction="none"
)
def _example(*questions, weight=1.0):
return DecisionLabelExample(questions=tuple(questions), source_weight=weight)
def _question(position, ids, target):
return DecisionLabelQuestion(position, tuple(ids), target)
def _single(position, ids, target, weight=1.0):
return _example(_question(position, ids, target), weight=weight)
def _dist(probs, candidates=None, other=None):
fields = {} if candidates is None else {"candidate_indices": candidates}
if other is not None:
fields["other_probability"] = other
return DistributionLabel(tuple(probs), **fields)
def _row(*values, grad=False, dtype=None):
return torch.tensor([[list(values)]], dtype=dtype, requires_grad=grad)
MIXED_EXAMPLES = [
_single(0, (0, 3), SetLabel((0, 1)), weight=1.75),
_example(
_question(1, (1, 2, 4), _dist((0.0, 0.4, 0.6))),
_question(2, (2, 5), HardLabel(0)),
weight=0.25,
),
]
def _loss(logits, examples, mask=ONE, **kwargs):
if not isinstance(examples, list):
examples = [examples]
return decision_label_loss(logits, examples, mask, **kwargs)
def _hidden_loss(hidden, head, examples, mask=ONE, **kwargs):
if not isinstance(examples, list):
examples = [examples]
kwargs.setdefault("linear_token_loss", _dense_linear_token_loss)
return decision_label_loss_from_hidden(hidden, head, examples, mask, **kwargs)
def _assert_fields_close(actual, expected, names=RESULT_FIELDS, **tolerance):
for name in names:
torch.testing.assert_close(
getattr(actual, name), getattr(expected, name), **tolerance
)
def _assert_finite_backward(result, *leaves):
assert torch.isfinite(result.loss)
result.loss.backward()
for leaf in leaves:
assert leaf.grad is not None and torch.isfinite(leaf.grad).all()
@pytest.mark.parametrize("mode", ["both", "full", "restricted"])
def test_per_example_components_average_to_aggregates_for_mixed_examples(mode):
torch.manual_seed(41)
hidden = torch.randn(2, 3, 4, dtype=torch.float64)
head = torch.nn.Linear(4, 7, dtype=torch.float64)
result = _hidden_loss(hidden, head, MIXED_EXAMPLES, MIXED_MASK, label_softmax=mode)
for aggregate, values in [
(result.loss, result.per_example_loss),
(result.restricted_loss, result.per_example_restricted_loss),
(result.full_vocab_loss, result.per_example_full_vocab_loss),
(result.brier_loss, result.per_example_brier_loss),
]:
assert values is not None
torch.testing.assert_close(torch.stack(values).mean(), aggregate)
assert all(not value.requires_grad for value in values)
@pytest.mark.parametrize("mode", ["both", "full", "restricted"])
def test_hidden_loss_accepts_accelerate_fp32_outputs_with_bf16_head(mode):
from accelerate.utils import convert_outputs_to_fp32
torch.manual_seed(19)
encoder = torch.nn.Linear(4, 4)
head = torch.nn.Linear(4, 3, dtype=torch.bfloat16)
@convert_outputs_to_fp32
def forward(values):
with torch.autocast("cpu", dtype=torch.bfloat16):
return {"last_hidden_state": encoder(values)}
hidden = forward(torch.ones(1, 1, 4))["last_hidden_state"]
hidden.retain_grad()
assert hidden.dtype == torch.float32
example = _single(0, (0, 1, 2), HardLabel(1))
seen = []
def callback(query_hidden, callback_head, targets):
seen.append(query_hidden.dtype)
return _dense_linear_token_loss(query_hidden, callback_head, targets)
result = _hidden_loss(
hidden, head, example, linear_token_loss=callback, label_softmax=mode
)
reference = _loss(
F.linear(hidden.to(head.weight.dtype), head.weight, head.bias).float(),
example,
label_softmax=mode,
)
torch.testing.assert_close(result.loss, reference.loss)
result.loss.backward()
assert seen == ([] if mode == "restricted" else [torch.bfloat16])
assert hidden.grad.dtype == torch.float32
assert torch.isfinite(hidden.grad).all() and torch.count_nonzero(hidden.grad)
assert encoder.weight.grad is not None
assert torch.isfinite(encoder.weight.grad).all()
assert head.weight.dtype == torch.bfloat16
@pytest.mark.parametrize(
("target", "mode"),
[
(HardLabel(1), "both"),
(_dist((0.0, 0.25, 0.75)), "both"),
(SetLabel((0, 2)), "both"),
(HardLabel(1), "full"),
(SetLabel((0, 2)), "restricted"),
],
)
def test_hidden_cce_loss_matches_dense_callback_and_backpropagates(target, mode):
torch.manual_seed(7)
hidden = torch.randn(2, 3, 4, dtype=torch.float64, requires_grad=True)
head = torch.nn.Linear(4, 7, bias=True, dtype=torch.float64)
examples = [
_single(0, (1, 3, 5), target, weight=2.0),
_example(
_question(1, (0, 2), HardLabel(0)),
_question(2, (1, 4, 6), _dist((0.0, 0.5, 0.5))),
weight=0.5,
),
]
calls = []
def callback(query_hidden, callback_head, targets):
calls.append((query_hidden.shape, targets.shape))
return _dense_linear_token_loss(query_hidden, callback_head, targets)
result = _hidden_loss(
hidden,
head,
examples,
MIXED_MASK,
linear_token_loss=callback,
label_softmax=mode,
brier_weight=0.3,
)
_assert_finite_backward(result, hidden, head.weight, head.bias)
expected_calls = (
[] if mode == "restricted" else [(torch.Size([3, 4]), torch.Size([3]))]
)
assert calls == expected_calls
if mode == "full":
torch.testing.assert_close(
result.restricted_loss, torch.zeros_like(result.restricted_loss)
)
def test_hidden_cce_helper_rejects_reduced_callback_output():
head = torch.nn.Linear(2, 3, bias=False, dtype=torch.float64)
with pytest.raises(ValueError, match="one loss per selected"):
_hidden_loss(
torch.zeros(1, 1, 2, dtype=torch.float64),
head,
_single(0, (0, 1), HardLabel(0)),
linear_token_loss=lambda values, linear, targets: _dense_linear_token_loss(
values, linear, targets
).mean(),
)
@pytest.mark.parametrize(
"target", [HardLabel(150), _dist((0.0,) * 150 + (1.0,))], ids=["hard", "dist"]
)
def test_151_candidate_final_index_has_normalized_loss_and_gradients(target):
hidden = torch.randn(1, 1, 4, dtype=torch.float64, requires_grad=True)
head = torch.nn.Linear(4, 151, bias=True, dtype=torch.float64)
logits = F.linear(hidden, head.weight, head.bias)
result = _loss(logits, _single(0, tuple(range(151)), target))
_assert_finite_backward(result, hidden, head.weight)
@pytest.mark.parametrize(
("examples", "mask", "kwargs", "fields", "dtype"),
[
(MIXED_EXAMPLES, MIXED_MASK, {}, RESULT_FIELDS, torch.float64),
(
MIXED_EXAMPLES,
MIXED_MASK,
{"full_ce_weighting": "dft"},
RESULT_FIELDS,
torch.float64,
),
(
_single(0, (0, 3), _dist((0.0, 1.0))),
ONE,
{**FULL, "hard_label_smoothing": 0.1},
("loss", "full_vocab_loss", "brier_loss"),
torch.float64,
),
(
_single(0, (0, 1), _dist((0.5,), (1,), 0.5)),
ONE,
{**FULL, "brier_weight": 0.0},
("loss", "full_vocab_loss"),
torch.float32,
),
],
ids=[
"ce-typed-unequal-questions",
"dft-typed-unequal-questions",
"one-hot-distribution-smoothing",
"sparse-distribution-residual",
],
)
def test_hidden_path_matches_dense_loss_and_gradients(
examples, mask, kwargs, fields, dtype
):
torch.manual_seed(11)
hidden = torch.randn(*mask.shape, 4, dtype=dtype, requires_grad=True)
head = torch.nn.Linear(4, 7, dtype=dtype)
dense_hidden = hidden.detach().clone().requires_grad_()
dense_head = torch.nn.Linear(4, 7, dtype=dtype)
dense_head.load_state_dict(head.state_dict())
cce = _hidden_loss(hidden, head, examples, mask, **kwargs)
logits = F.linear(dense_hidden, dense_head.weight, dense_head.bias)
dense = _loss(logits, examples, mask, **kwargs)
_assert_fields_close(cce, dense, fields)
cce.loss.backward()
dense.loss.backward()
torch.testing.assert_close(hidden.grad, dense_hidden.grad)
torch.testing.assert_close(head.weight.grad, dense_head.weight.grad)
torch.testing.assert_close(head.bias.grad, dense_head.bias.grad)
@pytest.mark.parametrize(
("examples", "mask", "variant"),
[
(
[_single(0, (0, 2, 3), HardLabel(1)), _single(1, (1, 3), SetLabel((0, 1)))],
torch.tensor([[True, False], [False, True]]),
"ce",
),
(
[
_example(
_question(0, (0, 1, 2), _dist((0.2, 0.3, 0.5))),
_question(1, (0, 2, 3), SetLabel((0, 2))),
),
_single(0, (1, 2, 3), _dist((0.5, 0.5, 0.0))),
],
torch.tensor([[True, True], [True, False]]),
"dft",
),
],
ids=["explicit-ce-is-default", "dft-leaves-distribution-and-set-unchanged"],
)
def test_full_ce_weighting_variants_preserve_loss_and_gradients(
examples, mask, variant
):
logits = torch.tensor(
[
[[0.4, -0.2, 1.3, 0.1], [0.8, -0.1, 0.2, 1.1]],
[[0.3, 0.9, -0.6, 0.2], [-0.4, 0.5, 1.0, 0.7]],
],
dtype=torch.float64,
)
baseline_logits = logits.clone().requires_grad_()
variant_logits = logits.clone().requires_grad_()
baseline = _loss(baseline_logits, examples, mask)
result = _loss(variant_logits, examples, mask, full_ce_weighting=variant)
baseline.loss.backward()
result.loss.backward()
torch.testing.assert_close(result.loss, baseline.loss, rtol=0, atol=0)
torch.testing.assert_close(
variant_logits.grad, baseline_logits.grad, rtol=0, atol=0
)
torch.testing.assert_close(
result.effective_full_vocab_loss, result.full_vocab_loss, rtol=0, atol=0
)
zero = torch.zeros_like(result.loss)
torch.testing.assert_close(result.full_vocab_dft_hard_weight_sum, zero)
torch.testing.assert_close(result.full_vocab_dft_hard_count, zero)
def test_dft_full_ce_uses_detached_target_probability_and_preserves_raw_metric():
logits = _row(0.3, -0.5, 1.1, grad=True, dtype=torch.float64)
expected_logits = logits.detach().clone().requires_grad_()
result = _loss(
logits,
_single(0, (0, 1, 2), HardLabel(2)),
full_ce_weighting="dft",
brier_weight=0,
**FULL,
)
full_ce = -F.log_softmax(logits[0, 0].float(), dim=0)[2]
probability = full_ce.detach().neg().exp()
torch.testing.assert_close(result.full_vocab_loss, full_ce)
torch.testing.assert_close(result.effective_full_vocab_loss, probability * full_ce)
torch.testing.assert_close(result.full_vocab_dft_hard_weight_sum, probability)
torch.testing.assert_close(
result.full_vocab_dft_hard_count, torch.ones_like(probability)
)
result.loss.backward()
(-F.softmax(expected_logits[0, 0].float(), dim=0)[2]).backward()
torch.testing.assert_close(logits.grad, expected_logits.grad)
def _nll(row, index, ids=None):
return -F.log_softmax(row if ids is None else row[list(ids)], dim=0)[index]
def _hard_components(row):
restricted, full = _nll(row, 1, (0, 3)), _nll(row, 3)
brier = (F.softmax(row[[0, 3]], dim=0) - torch.tensor([0.0, 1.0])).square().sum()
return {
"restricted_loss": restricted,
"full_vocab_loss": full,
"brier_loss": brier,
"loss": restricted + full + 0.5 * brier,
}
def _distribution_components(row):
target = torch.tensor([0.25, 0.75])
kl = torch.sum(target * (target.log() - F.log_softmax(row[[1, 2]], dim=0)))
full = -(target * F.log_softmax(row, dim=0)[[1, 2]]).sum()
return {"restricted_loss": kl, "full_vocab_loss": full, "loss": kl + full}
def _sparse_residual_components(row):
logprobs = F.log_softmax(row, dim=0)
full = -(torch.tensor([0.25, 0.5]) * logprobs[[1, 2]]).sum()
full = full - 0.25 * torch.log1p(-logprobs[[1, 2]].exp().sum())
return {"full_vocab_loss": full, "loss": full}
def _set_components(row):
mass = F.softmax(row[[0, 1, 2]], dim=0)[[0, 2]].sum()
full = _nll(row, 0)
return {
"restricted_loss": -mass.log(),
"full_vocab_loss": full,
"brier_loss": (1.0 - mass).square(),
"loss": -mass.log() + full + 0.25 * (1.0 - mass).square(),
}
ZERO = torch.zeros(())
@pytest.mark.parametrize(
("logits", "question", "kwargs", "expected"),
[
(
torch.tensor([[[30.0, -30.0, 20.0, -20.0], [0.2, 1.5, -0.4, 0.7]]]),
_question(1, (0, 3), HardLabel(1)),
{"brier_weight": 0.5},
_hard_components,
),
(
_row(0.3, -0.2, 1.1, 0.0),
_question(0, (1, 2), _dist((0.25, 0.75))),
{"brier_weight": 0.0},
_distribution_components,
),
(
_row(0.3, -0.2, 1.1, 0.0),
_question(0, (0, 1, 2), _dist((0.25, 0.5), (1, 2), 0.25)),
{**FULL, "brier_weight": 0.0},
_sparse_residual_components,
),
(
_row(0.1, 0.2, 1.5),
_question(0, (0, 1, 2), _dist((0.0, 0.0, 1.0))),
{"brier_weight": 0},
lambda row: {"restricted_loss": _nll(row, 2)},
),
(
_row(1.0, -1.0, 0.4, 2.0),
_question(0, (0, 1, 2), SetLabel((0, 2))),
{"brier_weight": 0.25},
_set_components,
),
(
_row(0.3, 1.2, -0.1),
_question(0, (0, 1, 2), SetLabel((0, 1, 2))),
{"label_softmax": "restricted", "brier_weight": 0},
lambda row: {"restricted_loss": ZERO, "brier_loss": ZERO},
),
(
_row(0.1, 1.2, -0.2),
_question(0, (0, 2), HardLabel(0)),
{**FULL, "brier_weight": 0},
lambda row: {"restricted_loss": ZERO, "full_vocab_loss": _nll(row, 0)},
),
],
ids=[
"hard-combines-restricted-full-and-brier-at-supervised-position",
"distribution-restricted-kl-and-full-soft-ce",
"sparse-distribution-full-ce-keeps-residual-other-mass",
"distribution-kl-accepts-zero-probability-alternatives",
"set-allowed-mass-and-detached-best-member-full-ce",
"set-mass-zero-when-every-allowed-label-correct",
"full-mode-excludes-restricted-primary-for-hard",
],
)
def test_single_question_loss_components(logits, question, kwargs, expected):
mask = torch.zeros(logits.shape[:2], dtype=torch.bool)
mask[0, question.position] = True
result = _loss(logits, _example(question), mask, **kwargs)
assert torch.isfinite(result.loss)
for name, value in expected(logits[0, question.position]).items():
torch.testing.assert_close(getattr(result, name), value)
@pytest.mark.parametrize(
("allowed_ids", "target", "expected_target", "gold", "hard"),
[
((0, 3), _dist((0.0, 1.0)), (0.05, 0.95), 3, (0.0, 1.0)),
(
(0, 1, 2, 3),
_dist((0.0, 0.0, 1.0, 0.0)),
(0.025, 0.025, 0.925, 0.025),
2,
(0.0, 0.0, 1.0, 0.0),
),
((0, 1, 2, 3), _dist((1.0,), (3,)), (0.025, 0.025, 0.025, 0.925), 3, None),
],
ids=["two-way", "four-way", "sparse-one-hot"],
)
def test_one_hot_distribution_smoothing_uses_allowed_support_for_ce_and_keeps_hard_brier(
allowed_ids, target, expected_target, gold, hard
):
logits = _row(0.2, 1.5, -0.4, 0.7, 1.1, grad=True)
example = _single(0, allowed_ids, target)
result = _loss(logits, example, brier_weight=0.5, hard_label_smoothing=0.1, **FULL)
unsmoothed = _loss(logits, example, brier_weight=0.5, **FULL)
result.loss.backward()
full_logprobs = F.log_softmax(logits.detach()[0, 0], dim=0)
torch.testing.assert_close(
result.full_vocab_loss,
-(torch.tensor(expected_target) * full_logprobs[list(allowed_ids)]).sum(),
)
if hard is not None:
restricted_probs = F.softmax(logits.detach()[0, 0, list(allowed_ids)], dim=0)
torch.testing.assert_close(
result.brier_loss, (restricted_probs - torch.tensor(hard)).square().sum()
)
assert not torch.allclose(result.loss, unsmoothed.loss)
assert logits.grad is not None
assert logits.grad[0, 0, gold] < 0
for token_id in set(range(logits.shape[-1])) - set(allowed_ids):
assert logits.grad[0, 0, token_id] > 0
@pytest.mark.parametrize(
("logits", "examples", "kwargs"),
[
(
_row(0.0, 100.0, 0.0, 0.0, 0.0, grad=True),
_single(0, (0, 1), _dist((0.5,), (1,), 0.5)),
{**FULL, "brier_weight": 0.0},
),
(
_row(1000.0, -1000.0, 999.0, grad=True),
_single(0, (0, 1, 2), SetLabel((1,))),
{"label_softmax": "restricted", "brier_weight": 0},
),
(
torch.tensor(
[[[0.1, 0.4, -0.3]], [[0.1, 0.4, -0.3]]],
dtype=torch.float64,
requires_grad=True,
),
[
_single(0, (0, 1, 2), HardLabel(1)),
_single(0, (0, 1, 2), HardLabel(1, smoothing=0.1)),
],
{**FULL, "hard_label_smoothing": 0.0},
),
],
ids=[
"sparse-residual-when-candidate-mass-rounds-to-one",
"set-mass-logsumexp-for-extremely-unlikely-correct-label",
"mixed-per-target-smoothing-with-global-smoothing-disabled",
],
)
def test_loss_and_gradients_stay_finite_at_numerical_edges(logits, examples, kwargs):
mask = torch.ones(logits.shape[:2], dtype=torch.bool)
_assert_finite_backward(_loss(logits, examples, mask, **kwargs), logits)
def test_distribution_full_ce_has_soft_target_gradient_and_ignores_zero_mass():
logits = _row(0.7, -0.3, 1.2, 0.1, grad=True)
example = _single(0, (0, 2, 3), _dist((0.25, 0.75, 0.0)))
result = _loss(logits, example, brier_weight=0.0, **FULL)
result.loss.backward()
target = torch.tensor([0.25, 0.75, 0.0])
expected = -(target * F.log_softmax(logits.detach()[0, 0], 0)[[0, 2, 3]]).sum()
torch.testing.assert_close(result.full_vocab_loss, expected)
assert logits.grad is not None
assert logits.grad[0, 0, 0] != 0
assert logits.grad[0, 0, 2] != 0
def test_question_then_example_average_applies_source_weights():
examples = [
_example(
_question(0, (0, 1), HardLabel(0)),
_question(1, (2, 3), HardLabel(1)),
weight=2.0,
),
_single(0, (0, 1, 2, 3), HardLabel(2), weight=0.5),
]
result = _loss(
torch.zeros(2, 2, 4),
examples,
torch.tensor([[True, True], [True, False]]),
label_softmax="restricted",
brier_weight=0,
)
expected = (2.0 * math.log(2) + 0.5 * math.log(4)) / 2
torch.testing.assert_close(result.loss, torch.tensor(expected))
def test_question_weights_scale_each_question_within_the_canvas_average():
# Two questions on one canvas: weight 3 on the first, 1 on the second.
examples = [
_example(
DecisionLabelQuestion(0, (0, 1), HardLabel(0), weight=3.0),
DecisionLabelQuestion(1, (0, 1, 2, 3), HardLabel(1)),
)
]
result = _loss(
torch.zeros(1, 2, 4),
examples,
torch.tensor([[True, True]]),
label_softmax="restricted",
brier_weight=0,
)
expected = (3.0 * math.log(2) + 1.0 * math.log(4)) / 2
torch.testing.assert_close(result.loss, torch.tensor(expected))
# Weight 1 everywhere is the plain question mean.
plain = _loss(
torch.zeros(1, 2, 4),
[
_example(
_question(0, (0, 1), HardLabel(0)),
_question(1, (0, 1, 2, 3), HardLabel(1)),
)
],
torch.tensor([[True, True]]),
label_softmax="restricted",
brier_weight=0,
)
torch.testing.assert_close(
plain.loss, torch.tensor((math.log(2) + math.log(4)) / 2)
)
def test_mapping_label_weight_is_read_and_validated():
from axolotl.integrations.diffusion_decision.loss import _question_weight
assert _question_weight({"kind": "hard", "gold_idx": 0}) == 1.0
assert _question_weight({"kind": "dist", "probs": [1.0], "weight": 0.25}) == 0.25
with pytest.raises(ValueError, match="label weight"):
_question_weight({"kind": "hard", "gold_idx": 0, "weight": 0})
with pytest.raises(ValueError, match="label weight"):
_question_weight({"kind": "hard", "gold_idx": 0, "weight": float("nan")})
def test_loss_only_backpropagates_through_supervised_label_positions():
logits = torch.randn(1, 3, 5, requires_grad=True)
mask = torch.tensor([[False, True, False]])
result = _loss(logits, _single(1, (1, 3), HardLabel(0)), mask, brier_weight=0)
result.loss.backward()
assert logits.grad is not None
assert not torch.count_nonzero(logits.grad[0, 0])
assert torch.count_nonzero(logits.grad[0, 1])
assert not torch.count_nonzero(logits.grad[0, 2])
@pytest.mark.parametrize(
("target", "error"),
[
({"kind": "hard", "gold_idx": -1}, "gold_idx"),
({"kind": "dist", "probs": [0.2, 0.2]}, "sum to one"),
({"kind": "set", "allowed_set": [0, 0]}, "unique"),
({"kind": "other"}, "unknown"),
],
)
def test_mapping_target_contracts(target, error):
with pytest.raises(ValueError, match=error):
label_target_from_mapping(target)
@pytest.mark.parametrize(
("example", "kwargs", "error", "match"),
[
(_single(1, (0, 1), _dist((0.5, 0.5))), {}, RuntimeError, "supervised"),
(_single(0, (0, 1), _dist((1.0,))), {}, ValueError, "align"),
(
_example(
_question(0, (0, 1), HardLabel(0)), _question(0, (1, 2), HardLabel(1))
),
{},
ValueError,
"distinct",
),
(
_single(0, (0, 1), _dist((0.75,), (1,), 0.25)),
{"label_softmax": "both"},
ValueError,
"requires label_softmax=full",
),
(_single(True, (0, 1), HardLabel(0)), {}, ValueError, None),
(_single(0, (True, 1), HardLabel(0)), {}, ValueError, None),
(_single(0, (0, 1), HardLabel(True)), {}, ValueError, None),
(_single(0, (0, 1), SetLabel((True,))), {}, ValueError, None),
],
ids=[
"unsupervised-label-position",
"misaligned-distribution-target",
"duplicate-question-positions",
"sparse-residual-needs-full-softmax",
"boolean-position",
"boolean-allowed-id",
"boolean-gold-index",
"boolean-set-member",
],
)
def test_loss_rejects_invalid_examples(example, kwargs, error, match):
with pytest.raises(error, match=match):
_loss(torch.zeros(1, 2, 4), example, torch.tensor([[True, False]]), **kwargs)
@pytest.mark.parametrize(
("left", "right", "fields", "exact"),
[
(
(_single(0, (0, 3), _dist((0.25, 0.75))), {**FULL, "brier_weight": 0.1}),
(
_single(0, (0, 3), _dist((0.25, 0.75))),
{**FULL, "brier_weight": 0.1, "hard_label_smoothing": 0.1},
),
("loss",),
False,
),
(
(_single(0, (0, 3), _dist((0.0, 1.0))), FULL),
(
_single(0, (0, 3), _dist((0.0, 1.0))),
{**FULL, "hard_label_smoothing": 0.0},
),
("loss",),
False,
),
(
(
_single(0, (0, 2, 4), HardLabel(1)),
{**FULL, "brier_weight": 0.1, "hard_label_smoothing": 0.1},
),
(
_single(0, (0, 2, 4), HardLabel(1, smoothing=0.1)),
{**FULL, "brier_weight": 0.1, "hard_label_smoothing": 0.0},
),
("loss", "full_vocab_loss", "brier_loss"),
True,
),
(
(_single(0, (1, 2), _dist((0.0, 1.0))), {"brier_weight": 0.0}),
(_single(0, (1, 2), HardLabel(1)), {"brier_weight": 0.0}),
("loss", "full_vocab_loss"),
False,
),
],
ids=[
"hard-smoothing-leaves-soft-distribution-unchanged",
"zero-hard-smoothing-preserves-one-hot-distribution",
"per-target-smoothing-matches-global-and-keeps-one-hot-brier",
"one-hot-distribution-full-ce-matches-hard-label",
],
)
def test_equivalent_objectives_produce_the_same_loss(left, right, fields, exact):
logits = _row(0.2, -0.7, 1.3, 0.4, 0.9, dtype=torch.float64)
tolerance = {"rtol": 0, "atol": 0} if exact else {}
_assert_fields_close(
_loss(logits, left[0], **left[1]),
_loss(logits, right[0], **right[1]),
fields,
**tolerance,
)
tests/integrations/diffusion_decision/test_manifest.py
+72
View File
@@ -0,0 +1,72 @@
"""Decision manifest compatibility and runtime-cost coverage."""
import json
from types import SimpleNamespace
from axolotl.integrations.diffusion_decision.manifest import (
DecisionManifest,
build_decision_manifest,
)
from axolotl.utils.dict import DictDefault
class _ModelWithoutParameterTraversal:
config = SimpleNamespace(mask_token_id=100, _commit_hash="resolved-revision")
def named_parameters(self):
raise AssertionError("manifest generation must not traverse model parameters")
def _manifest_config():
return DictDefault(
{
"base_model": "nvidia/Nemotron-Labs-Diffusion-3B",
"revision_of_model": "requested-revision",
"model_config_type": "nemotron_labs_diffusion",
"diffusion_lm": {"canvas_width": 128, "mask_token_id": 100},
"diffusion_decision": {
"layout": "prompt_slots",
"latent": {"mode": "none", "num_slots": 0},
},
}
)
def test_manifest_build_does_not_traverse_model_parameters():
manifest = build_decision_manifest(
_manifest_config(), model=_ModelWithoutParameterTraversal()
)
assert manifest.initial_common_adapter_sha256 is None
assert manifest.initial_common_adapter_fingerprint_reason is None
assert manifest.initial_common_adapter_devices == []
assert manifest.initial_common_adapter_dtypes == []
assert manifest.initial_common_adapter_tensor_count == 0
def test_manifest_parses_legacy_initial_adapter_fingerprint_fields():
manifest = build_decision_manifest(_manifest_config())
legacy = manifest.model_dump(mode="json") | {
"initial_common_adapter_sha256": "a" * 64,
"initial_common_adapter_fingerprint_reason": None,
"initial_common_adapter_devices": ["cuda:0"],
"initial_common_adapter_dtypes": ["torch.bfloat16"],
"initial_common_adapter_tensor_count": 2,
}
restored = DecisionManifest.model_validate_json(json.dumps(legacy))
assert restored.initial_common_adapter_sha256 == "a" * 64
assert restored.initial_common_adapter_devices == ["cuda:0"]
assert restored.initial_common_adapter_dtypes == ["torch.bfloat16"]
assert restored.initial_common_adapter_tensor_count == 2
def test_manifest_drops_retired_read_protocol_fields():
manifest = build_decision_manifest(_manifest_config())
legacy = manifest.model_dump(mode="json")
legacy["noise_read_protocol"] |= {"reader": "hf", "noise_draws": 2}
restored = DecisionManifest.model_validate_json(json.dumps(legacy))
assert restored == manifest
tests/integrations/diffusion_decision/test_mixture.py
+251
View File
@@ -0,0 +1,251 @@
from collections import Counter
import pytest
from axolotl.integrations.diffusion_decision.mixture import (
DeterministicMixtureSampler,
prepare_source_pools,
)
def _record(source, group, state, questions=1):
return {
"source": source,
"group": group,
"state": state,
"questions": {str(index): {} for index in range(questions)},
}
def test_prepare_source_pools_decontaminates_and_caps_each_source_deterministically():
training = [
_record("alpha", "alpha-0", "keep-a"),
_record("alpha", "alpha-1", "keep-b"),
_record("alpha", "alpha-2", "overlap"),
_record("beta", "beta-0", "keep-c"),
]
evaluation = [_record("eval", "eval-0", "overlap")]
first = prepare_source_pools(
training, evaluation, max_examples_per_source={"alpha": 1}, seed=7
)
second = prepare_source_pools(
training, evaluation, max_examples_per_source={"alpha": 1}, seed=7
)
assert first.records == second.records
assert set(first.records) == {"alpha", "beta"}
assert len(first.records["alpha"]) == 1
assert first.counts == {
"input": 4,
"state_overlap": 1,
"family_overlap": 0,
"kept": 3,
"capped": 1,
"sources": 2,
}
assert all(
record["state"] != "overlap"
for pool in first.records.values()
for record in pool
)
def test_prepare_source_pools_removes_training_families_held_out_for_evaluation():
training = [
_record("alpha", "shared-family", "train-only"),
_record("beta", "beta-family", "keep"),
]
evaluation = [_record("alpha", "shared-family", "eval-only")]
prepared = prepare_source_pools(training, evaluation)
assert set(prepared.records) == {"beta"}
assert prepared.counts["family_overlap"] == 1
def test_explicit_weight_draw_shares_match_within_two_percent_over_ten_thousand():
sampler = DeterministicMixtureSampler(
{"a": tuple(range(3)), "b": tuple(range(5)), "c": tuple(range(7))},
weights={"a": 0.5, "b": 0.3, "c": 0.2},
seed=19,
)
counts = Counter(sample.source for sample in sampler.draw(10_000))
for source, expected in sampler.probabilities.items():
assert abs(counts[source] / 10_000 - expected) <= 0.02
def test_temperature_uses_capped_source_sizes():
sampler = DeterministicMixtureSampler(
{"small": tuple(range(4)), "large": tuple(range(16))}, temperature=0.5
)
assert sampler.probabilities == {
"large": pytest.approx(2 / 3),
"small": pytest.approx(1 / 3),
}
def test_probability_normalization_handles_finite_extreme_explicit_weights():
sampler = DeterministicMixtureSampler(
{"a": (1,), "b": (2,)}, weights={"a": 1e308, "b": 1e308}
)
assert sampler.probabilities == {"a": 0.5, "b": 0.5}
def test_temperature_zero_is_uniform_and_large_finite_temperature_is_stable():
uniform = DeterministicMixtureSampler(
{"small": (1,), "large": tuple(range(16))}, temperature=0
)
stable = DeterministicMixtureSampler(
{"small": (1,), "large": (2,)}, temperature=1000
)
assert uniform.probabilities == {"large": 0.5, "small": 0.5}
assert all(0 < probability < 1 for probability in stable.probabilities.values())
@pytest.mark.parametrize(
("pools", "weights", "temperature"),
[
(
{"small": (1,), "large": tuple(range(2))},
{"small": 1e-308, "large": 1e308},
None,
),
({"small": (1,), "large": tuple(range(2))}, None, 2000),
({"small": (1,), "large": tuple(range(16))}, None, 1e308),
],
)
def test_probability_construction_rejects_source_underflow_or_overflow(
pools, weights, temperature
):
with pytest.raises(ValueError, match="underflowed|nonfinite"):
DeterministicMixtureSampler(pools, weights=weights, temperature=temperature)
def test_stratified_batches_include_sources_with_at_least_one_expected_slot():
sampler = DeterministicMixtureSampler(
{source: tuple(range(3)) for source in ("a", "b", "c", "d")},
weights={"a": 0.5, "b": 0.25, "c": 0.125, "d": 0.125},
seed=3,
)
batches = list(sampler.batches(8, 30, stratified=True))
assert all(len(batch) == 8 for batch in batches)
assert all(
{sample.source for sample in batch} == {"a", "b", "c", "d"} for batch in batches
)
assert all(
Counter(sample.source for sample in batch) == {"a": 4, "b": 2, "c": 1, "d": 1}
for batch in batches
)
def test_fractional_stratification_converges_to_source_probabilities():
sampler = DeterministicMixtureSampler(
{source: tuple(range(3)) for source in ("a", "b", "c")},
weights={"a": 0.6, "b": 0.25, "c": 0.15},
seed=3,
)
counts = Counter(
sample.source
for batch in sampler.batches(4, 2_500, stratified=True)
for sample in batch
)
for source, expected in sampler.probabilities.items():
assert abs(counts[source] / 10_000 - expected) <= 0.02
def test_seed_reproduces_draws_and_stratified_batches():
pools = {"a": tuple(range(4)), "b": tuple(range(6))}
first = DeterministicMixtureSampler(pools, weights={"a": 0.4, "b": 0.6}, seed=11)
second = DeterministicMixtureSampler(pools, weights={"a": 0.4, "b": 0.6}, seed=11)
assert first.draw(40) == second.draw(40)
assert list(first.batches(5, 20)) == list(second.batches(5, 20))
def test_epoch_batches_cover_every_capped_source_record_before_stopping():
sampler = DeterministicMixtureSampler(
{"small": tuple(range(2)), "large": tuple(range(7))},
weights={"small": 0.5, "large": 0.5},
seed=11,
)
batches = list(sampler.epoch_batches(4))
records = {
source: {
sample.record
for batch in batches
for sample in batch
if sample.source == source
}
for source in ("small", "large")
}
assert all(len(batch) == 4 for batch in batches)
assert records == {"small": {0, 1}, "large": set(range(7))}
def test_non_stratified_epoch_is_seeded_and_covers_every_source_record():
pools = {"a": tuple(range(2)), "b": tuple(range(5))}
first = DeterministicMixtureSampler(pools, weights={"a": 0.2, "b": 0.8}, seed=7)
second = DeterministicMixtureSampler(pools, weights={"a": 0.2, "b": 0.8}, seed=7)
batches = list(first.epoch_batches(3, stratified=False))
assert batches == list(second.epoch_batches(3, stratified=False))
assert all(len(batch) == 3 for batch in batches)
assert {
(sample.source, sample.record) for batch in batches for sample in batch
} >= {("a", 0), ("a", 1), *{("b", value) for value in range(5)}}
@pytest.mark.parametrize(
("pools", "weights", "temperature", "error"),
[
({}, None, None, "at least one"),
({"a": ()}, None, None, "empty"),
({"a": (1,)}, {"missing": 1.0}, None, "every nonempty"),
({"a": (1,)}, {"a": 1.0}, 0.5, "or temperature"),
({"a": (1,)}, {"a": 0.0}, None, "positive"),
],
)
def test_sampler_rejects_empty_or_invalid_source_definitions(
pools, weights, temperature, error
):
with pytest.raises(ValueError, match=error):
DeterministicMixtureSampler(pools, weights=weights, temperature=temperature)
def test_examples_remain_whole_records_when_sampled():
record = _record("alpha", "family", "state", questions=3)
sampler = DeterministicMixtureSampler({"alpha": (record,)}, seed=2)
sample = sampler.draw(1)[0]
assert sample.source == "alpha"
assert sample.record is record
assert len(sample.record["questions"]) == 3
def test_prepare_rejects_cap_for_source_removed_by_decontamination():
training = [_record("alpha", "alpha", "overlap")]
evaluation = [_record("eval", "eval", "overlap")]
with pytest.raises(ValueError, match="unavailable"):
prepare_source_pools(training, evaluation, max_examples_per_source={"alpha": 1})
def test_prepare_rejects_when_decontamination_removes_every_source():
training = [_record("alpha", "alpha", "overlap")]
evaluation = [_record("eval", "eval", "overlap")]
with pytest.raises(ValueError, match="no nonempty"):
prepare_source_pools(training, evaluation)
tests/integrations/diffusion_decision/test_nemotron_free_varlen.py
+196
View File
@@ -0,0 +1,196 @@
"""CUDA smoke coverage for native Nemotron free-slot decision recurrence."""
from __future__ import annotations
import pytest
import torch
from peft import LoraConfig, TaskType, get_peft_model
from transformers.dynamic_module_utils import get_class_from_dynamic_module
from axolotl.core.trainers.diffusion_lm.backends.full_sequence import (
FullSequenceBackend,
)
from axolotl.integrations.diffusion_decision.loss import (
DecisionLabelExample,
DecisionLabelQuestion,
HardLabel,
)
from axolotl.model_support import DiffusionLayout, LogitAlignment
from axolotl.model_support.nemotron_diffusion.compat import (
resolve_nemotron_model_class,
)
from tests.integrations.diffusion_decision.helpers import (
MultistepTrainerHarness,
trainer_spec,
)
from tests.native_source_fixtures import native_source_fixture_path
def _model(device: torch.device):
source = native_source_fixture_path("nemotron")
if source is None:
pytest.skip("native Nemotron source fixture unavailable")
config_class = get_class_from_dynamic_module(
"configuration_nemotron_labs_diffusion.NemotronLabsDiffusionConfig",
str(source),
local_files_only=True,
)
config = config_class(
vocab_size=128,
hidden_size=256,
intermediate_size=512,
num_hidden_layers=1,
num_attention_heads=4,
num_key_value_heads=2,
head_dim=64,
max_position_embeddings=64,
mask_token_id=100,
dlm_paradigm="bidirectional",
rope_parameters={
"llama_4_scaling_beta": 1.0,
"original_max_position_embeddings": 1,
},
use_cache=False,
)
config._attn_implementation = "eager"
base = resolve_nemotron_model_class(str(source))(config).to(
device=device, dtype=torch.bfloat16
)
traces: list[tuple[torch.Tensor, torch.Tensor]] = []
original_forward = base.forward
def traced_forward(*args, **kwargs):
result = original_forward(*args, **kwargs)
traces.append(
(kwargs["input_ids"].detach().clone(), result.logits.detach().clone())
)
return result
base.forward = traced_forward
model = get_peft_model(
base,
LoraConfig(
task_type=TaskType.CAUSAL_LM,
target_modules=["q_proj"],
r=2,
lora_alpha=2,
),
).train()
return model, traces, source
def _inputs(device: torch.device) -> dict[str, object]:
return {
"input_ids": torch.tensor([[1, 7, 11, 8, 2, 3, 9, 12, 10, 4]], device=device),
"document_ids": torch.tensor([[0, 0, 0, 0, 0, 1, 1, 1, 1, 1]], device=device),
"semantic_validity": torch.ones((1, 10), dtype=torch.bool, device=device),
"position_ids": torch.tensor([[0, 1, 2, 3, 4, 0, 1, 2, 3, 4]], device=device),
"canvas_corruptible_mask": torch.tensor(
[[False, False, True, False, False, False, False, True, False, False]],
device=device,
),
"canvas_input_pinned_mask": torch.tensor(
[[False, False, False, True, True, True, False, False, True, True]],
device=device,
),
"canvas_update_mask": torch.zeros((1, 10), dtype=torch.bool, device=device),
"decision_slot_mask": torch.tensor(
[[False, True, False, True, False, False, True, False, False, False]],
device=device,
),
"canvas_loss_mask": torch.tensor(
[[False, False, True, False, False, False, False, True, False, False]],
device=device,
),
"decision_examples": (
DecisionLabelExample(
questions=(DecisionLabelQuestion(0, (1, 2, 3), HardLabel(0)),)
),
DecisionLabelExample(
questions=(DecisionLabelQuestion(0, (4, 5, 6), HardLabel(1)),)
),
),
"decision_question_mask": torch.tensor([[True], [True]], device=device),
"decision_supervision_mask": torch.tensor([[True], [True]], device=device),
"decision_label_rows": torch.tensor([[0], [0]], device=device),
"decision_label_positions": torch.tensor([[2], [7]], device=device),
}
@pytest.mark.parametrize("steps", [2, 3])
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
def test_native_nemotron_varlen_free_slots_evolve_without_slot_loss(monkeypatch, steps):
from axolotl.core.trainers.diffusion_lm import varlen
device = torch.device("cuda")
torch.manual_seed(314159)
torch.cuda.reset_peak_memory_stats(device)
kernel_calls = []
original_kernel = varlen.varlen_attn
def traced_kernel(*args, **kwargs):
kernel_calls.append(args[3].detach().clone())
return original_kernel(*args, **kwargs)
monkeypatch.setattr(varlen, "varlen_attn", traced_kernel)
model, traces, source = _model(device)
trainer = MultistepTrainerHarness(
trainer_spec(DiffusionLayout.FULL_SEQUENCE, LogitAlignment.ALIGNED),
k_max=steps,
sampled_steps=steps,
decision={
"latent": {
"mode": "free",
"num_slots": 1,
"free_update_policy": "argmax",
}
},
)
trainer._full_sequence_backend = lambda: FullSequenceBackend(
mask_token_id=100, attention_backend="varlen"
)
inputs = _inputs(device)
loss = trainer.compute_loss(model, inputs)
loss.backward()
torch.cuda.synchronize(device)
slots = torch.tensor([1, 6], device=device)
pinned_slots = torch.tensor([3], device=device)
held = torch.tensor([0, 2, 3, 4, 5, 7, 8, 9], device=device)
metadata = FullSequenceBackend(mask_token_id=100, attention_backend="varlen").pack(
inputs["input_ids"], inputs["document_ids"], inputs["semantic_validity"]
)["diffusion_varlen"]
assert metadata is not None
assert metadata.cu_seqlens.tolist() == [0, 5, 10]
assert len(kernel_calls) == steps
assert all(cu_seqlens.tolist() == [0, 5, 10] for cu_seqlens in kernel_calls)
assert len(traces) == steps
assert all(
torch.equal(state[:, held], traces[0][0][:, held]) for state, _ in traces
)
assert not torch.any(inputs["decision_slot_mask"] & inputs["canvas_loss_mask"])
assert all(
torch.equal(state[:, pinned_slots], traces[0][0][:, pinned_slots])
for state, _ in traces
)
for (_state, logits), (next_state, _) in zip(traces[:-1], traces[1:], strict=True):
torch.testing.assert_close(next_state[:, slots], logits.argmax(-1)[:, slots])
assert torch.any(traces[1][0][:, slots] != traces[0][0][:, slots])
lora_grads = [
parameter.grad
for name, parameter in model.named_parameters()
if ".lora_" in name and parameter.requires_grad
]
assert torch.isfinite(loss)
assert lora_grads
assert all(
gradient is not None and torch.isfinite(gradient).all()
for gradient in lora_grads
)
assert any(torch.count_nonzero(gradient) for gradient in lora_grads)
print(
"native_nemotron_free_varlen"
f" steps={steps} fixture={source} torch={torch.__version__}"
f" peak_bytes={torch.cuda.max_memory_allocated(device)}"
)
tests/integrations/diffusion_decision/test_permutation.py
+129
View File
@@ -0,0 +1,129 @@
"""Option permutations preserve the meanings of all typed targets."""
from copy import deepcopy
import pytest
from axolotl.integrations.diffusion_decision.permutation import permute_record
def record(target):
return {
"source": "fixture",
"group": "group",
"id": "row",
"state": {"seen": "objects"},
"questions": {
"choice": {
"type": "choice",
"options": [
{"name": name, "description": f"description of {name}"}
for name in ("apple", "pear", "plum", "orange")
],
},
"binary": {"type": "noul"},
},
"labels": {"choice": target, "binary": {"kind": "hard", "gold_idx": 1}},
}
@pytest.mark.parametrize(
"target",
[
{"kind": "hard", "gold_idx": 2},
{"kind": "dist", "probs": [0.1, 0.2, 0.3, 0.4]},
{
"kind": "dist",
"candidate_indices": [0, 2],
"candidate_ids": ["apple", "plum"],
"probs": [0.1, 0.6],
"other_probability": 0.3,
},
{"kind": "set", "allowed_set": [0, 2]},
],
)
def test_semantic_targets_preserved_for_every_seed(target):
original = record(target)
untouched = deepcopy(original)
orders = set()
question_orders = set()
for seed in range(20):
result = permute_record(original, seed=seed)
assert result == permute_record(original, seed=seed)
options = result["questions"]["choice"]["options"]
names = [option["name"] for option in options]
assert all(
option["description"] == f"description of {option['name']}"
for option in options
)
remapped = result["labels"]["choice"]
if target["kind"] == "hard":
assert names[remapped["gold_idx"]] == "plum"
elif target["kind"] == "dist":
if "candidate_indices" in target:
assert {
names[index]: probability
for index, probability in zip(
remapped["candidate_indices"], remapped["probs"], strict=True
)
} == {"apple": 0.1, "plum": 0.6}
assert remapped["other_probability"] == 0.3
assert {
names[index]: identity
for index, identity in zip(
remapped["candidate_indices"],
remapped["candidate_ids"],
strict=True,
)
} == {"apple": "apple", "plum": "plum"}
else:
assert dict(zip(names, remapped["probs"], strict=True)) == {
"apple": 0.1,
"pear": 0.2,
"plum": 0.3,
"orange": 0.4,
}
else:
assert {names[index] for index in remapped["allowed_set"]} == {
"apple",
"plum",
}
assert result["labels"]["binary"] == untouched["labels"]["binary"]
assert result["state"] == untouched["state"]
orders.add(tuple(names))
question_orders.add(tuple(result["questions"]))
assert original == untouched
assert len(orders) > 1
assert len(question_orders) == 2
def test_question_input_order_does_not_change_seeded_result():
first = record({"kind": "hard", "gold_idx": 0})
second = deepcopy(first)
second["questions"] = dict(reversed(list(second["questions"].items())))
assert permute_record(first, seed=5) == permute_record(second, seed=5)
def test_expanded_codebook_is_preserved_during_permutation():
names = list("ABCDEFGHIJKLMNOPQRSTUVWXYZab")
original = {
"source": "fixture",
"group": "group",
"id": "expanded-choice",
"state": {"seen": "objects"},
"questions": {
"choice": {
"type": "choice",
"options": [
{"name": name, "description": f"description of {name}"}
for name in names
],
}
},
"labels": {"choice": {"kind": "dist", "probs": [1 / 28] * 28}},
}
result = permute_record(original, seed=7, codebook="expanded52")
assert len(result["questions"]["choice"]["options"]) == 28
assert result["labels"]["choice"]["probs"] == [1 / 28] * 28
tests/integrations/diffusion_decision/test_plugin_config.py
+428
View File
@@ -0,0 +1,428 @@
"""Decision plugin configuration and registration tests."""
from collections import OrderedDict
import pytest
from pydantic import ValidationError
from axolotl.integrations.diffusion_decision.args import (
DiffusionDecisionArgs,
DiffusionDecisionConfig,
)
from axolotl.integrations.diffusion_decision.plugin import DiffusionDecisionPlugin
PLUGIN_PATH = "axolotl.integrations.diffusion_decision.DiffusionDecisionPlugin"
def test_plugin_exposes_optional_nested_config():
assert (
DiffusionDecisionPlugin().get_input_args()
== "axolotl.integrations.diffusion_decision.args.DiffusionDecisionArgs"
)
assert DiffusionDecisionArgs().diffusion_decision is None
@pytest.mark.parametrize("layout", ("thought_block", "prompt_slots"))
def test_config_preserves_both_decision_layouts(layout):
config = DiffusionDecisionConfig.model_validate(
{
"layout": layout,
"read_fraction": 0.4,
"mixture": {"weights": {"one": 1.0}},
"labels": {"label_softmax": "both", "brier_weight": 0.1},
"latent": {"mode": "learned", "num_slots": 8},
}
)
assert config.layout == layout
assert config.latent.mode == "learned"
@pytest.mark.parametrize(
"block,message",
[
({"latent": {"mode": "none", "num_slots": 1}}, "mode=none"),
({"latent": {"token_ids": [-1]}}, "nonnegative"),
({"carry": {"enabled": True}}, "carry is not wired"),
({"mixture": {"weights": {"source": 0.0}}}, "finite and positive"),
(
{"mixture": {"weights": {"source": 1.0}, "temperature": 0.5}},
"explicit weights or temperature",
),
({"mixture": {"max_examples_per_source": {"source": 0}}}, "positive caps"),
({"labels": {"brier_weight": float("inf")}}, "finite"),
({"read_fraction": float("inf")}, "finite"),
],
)
def test_config_rejects_invalid_semantic_controls(block, message):
with pytest.raises(ValidationError, match=message):
DiffusionDecisionConfig.model_validate(block)
def test_config_accepts_per_source_caps_without_numeric_union_comparison():
config = DiffusionDecisionConfig.model_validate(
{"mixture": {"max_examples_per_source": {"source": 2}}}
)
assert config.mixture.max_examples_per_source == {"source": 2}
def test_labels_default_to_ce_weighting_and_accept_dft_only():
assert DiffusionDecisionConfig().labels.full_ce_weighting == "ce"
assert (
DiffusionDecisionConfig.model_validate(
{"labels": {"full_ce_weighting": "dft"}}
).labels.full_ce_weighting
== "dft"
)
with pytest.raises(ValidationError):
DiffusionDecisionConfig.model_validate(
{"labels": {"full_ce_weighting": "other"}}
)
def test_labels_accept_hard_label_smoothing_only_inside_unit_interval():
assert (
DiffusionDecisionConfig.model_validate(
{"labels": {"hard_label_smoothing": 0.1}}
).labels.hard_label_smoothing
== 0.1
)
with pytest.raises(ValidationError):
DiffusionDecisionConfig.model_validate(
{"labels": {"hard_label_smoothing": 1.0}}
)
with pytest.raises(ValidationError, match="requires full_ce_weighting=ce"):
DiffusionDecisionConfig.model_validate(
{"labels": {"hard_label_smoothing": 0.1, "full_ce_weighting": "dft"}}
)
def test_plugin_input_args_merge_into_config_schema(min_base_cfg, monkeypatch):
from axolotl.integrations.base import PluginManager
from axolotl.integrations.config import merge_input_args
manager = PluginManager.get_instance()
monkeypatch.setattr(
manager,
"plugins",
OrderedDict({PLUGIN_PATH: DiffusionDecisionPlugin()}),
)
_capabilities, config_cls = merge_input_args()
config = config_cls(
**(
min_base_cfg
| {
"plugins": [PLUGIN_PATH],
"diffusion_decision": {
"layout": "prompt_slots",
"latent": {"mode": "prompt", "num_slots": 4},
},
}
)
)
assert config.diffusion_decision.layout == "prompt_slots"
assert config.diffusion_decision.latent.num_slots == 4
@pytest.mark.parametrize("layout", ("thought_block", "prompt_slots"))
def test_native_registration_is_layout_agnostic(layout):
plugin = DiffusionDecisionPlugin()
plugin.register(
{
"diffusion_lm": {"from_causal_lm": False},
"diffusion_decision": {"layout": layout},
}
)
def test_registration_rejects_legacy_diffusion_compatibility():
with pytest.raises(ValueError, match="native diffusion_lm"):
DiffusionDecisionPlugin().register({"diffusion_decision": {}})
with pytest.raises(ValueError, match="from_causal_lm"):
DiffusionDecisionPlugin().register(
{
"diffusion_lm": {"from_causal_lm": True},
"diffusion_decision": {},
}
)
def test_registration_allows_sample_packing_with_top_level_budget():
DiffusionDecisionPlugin().register(
{
"sample_packing": True,
"micro_batch_size": 2,
"sequence_len": 2048,
"diffusion_lm": {"from_causal_lm": False},
"diffusion_decision": {},
}
)
def test_sample_packing_passes_sequence_len_times_micro_batch_capacity():
from types import SimpleNamespace
from axolotl.integrations.diffusion_decision.training_collator import (
decision_collator_for_config,
)
from axolotl.utils.dict import DictDefault
_collator, kwargs = decision_collator_for_config(
DictDefault(
{
"model_config_type": "nemotron_labs_diffusion",
"tokenizer": SimpleNamespace(pad_token_id=0),
"sample_packing": True,
"micro_batch_size": 3,
"sequence_len": 128,
"diffusion_lm": {"from_causal_lm": False},
}
)
)
assert kwargs["physical_payload_capacity"] == 384
def test_eval_packing_uses_the_training_physical_budget():
from types import SimpleNamespace
from axolotl.integrations.diffusion_decision.training_collator import (
decision_collator_for_config,
)
from axolotl.utils.dict import DictDefault
_collator, kwargs = decision_collator_for_config(
DictDefault(
{
"model_config_type": "nemotron_labs_diffusion",
"tokenizer": SimpleNamespace(pad_token_id=0),
"sample_packing": True,
"eval_sample_packing": True,
"micro_batch_size": 16,
"eval_batch_size": 8,
"sequence_len": 128,
"diffusion_lm": {"from_causal_lm": False},
}
),
is_eval=True,
)
assert kwargs["physical_payload_capacity"] == 2048
def test_fixed_logical_eval_uses_the_eval_physical_budget():
from types import SimpleNamespace
from axolotl.integrations.diffusion_decision.training_collator import (
decision_collator_for_config,
)
from axolotl.utils.dict import DictDefault
_collator, kwargs = decision_collator_for_config(
DictDefault(
{
"model_config_type": "nemotron_labs_diffusion",
"tokenizer": SimpleNamespace(pad_token_id=0),
"sample_packing": False,
"batch_flattening": True,
"micro_batch_size": 16,
"eval_batch_size": 8,
"sequence_len": 128,
"diffusion_lm": {"from_causal_lm": False},
}
),
is_eval=True,
)
assert kwargs["physical_payload_capacity"] == 1024
def test_registration_rejects_explicit_unflattened_microbatches():
config = {
"diffusion_lm": {"from_causal_lm": False},
"diffusion_decision": {},
"sample_packing": False,
"eval_sample_packing": True,
"micro_batch_size": 8,
}
config["batch_flattening"] = False
with pytest.raises(ValueError, match="require batch_flattening: true"):
DiffusionDecisionPlugin().register(config)
def test_fixed_logical_batch_requires_capacity():
config = {
"diffusion_lm": {
"from_causal_lm": False,
},
"diffusion_decision": {},
"batch_flattening": True,
"sample_packing": False,
"micro_batch_size": 8,
"sequence_len": 2048,
}
DiffusionDecisionPlugin().register(config)
def test_fixed_logical_batch_passes_physical_capacity_to_decision_collator():
from types import SimpleNamespace
from axolotl.integrations.diffusion_decision.training_collator import (
decision_collator_for_config,
)
from axolotl.utils.dict import DictDefault
_collator, kwargs = decision_collator_for_config(
DictDefault(
{
"model_config_type": "nemotron_labs_diffusion",
"tokenizer": SimpleNamespace(pad_token_id=0),
"sample_packing": False,
"batch_flattening": True,
"micro_batch_size": 8,
"sequence_len": 2048,
"diffusion_lm": {
"from_causal_lm": False,
"canvas_width": 128,
},
}
)
)
assert kwargs["physical_payload_capacity"] == 16384
def test_native_nemotron_plugin_routes_through_causal_builder(
min_base_cfg, monkeypatch
):
"""Exercise PluginManager and the causal builder with the resolved native profile."""
from types import SimpleNamespace
from axolotl.core.builders.causal import HFCausalTrainerBuilder
from axolotl.integrations.base import PluginManager
from axolotl.integrations.diffusion_decision.trainer import (
DiffusionDecisionTrainer,
)
from axolotl.integrations.diffusion_decision.training_collator import (
DecisionTrainingCollator,
)
from axolotl.model_support import (
DiffusionLayout,
MaskTokenPolicy,
get_model_support_for_cfg,
resolve_model_support,
)
from axolotl.utils.config import validate_config
from axolotl.utils.dict import DictDefault
manager = PluginManager.get_instance()
monkeypatch.setattr(
manager,
"plugins",
OrderedDict({PLUGIN_PATH: DiffusionDecisionPlugin()}),
)
cfg = validate_config(
DictDefault(
min_base_cfg
| {
"plugins": [PLUGIN_PATH],
"attn_implementation": "varlen",
"diffusion_lm": {
"from_causal_lm": False,
"canvas_width": 128,
"mask_token_id": 100,
},
"diffusion_decision": {"layout": "thought_block"},
}
)
)
cfg.model_config_type = "nemotron_labs_diffusion"
cfg.tokenizer = SimpleNamespace(pad_token_id=0)
profile = resolve_model_support(get_model_support_for_cfg(cfg))
assert profile is not None
assert profile.diffusion is not None
assert profile.diffusion.layout is DiffusionLayout.FULL_SEQUENCE
assert profile.diffusion.mask_token_policy is MaskTokenPolicy.MODEL
assert cfg.diffusion_lm.mask_token_id == 100
builder = HFCausalTrainerBuilder.__new__(HFCausalTrainerBuilder)
builder.cfg = cfg
builder.tokenizer = cfg.tokenizer
assert builder._get_trainer_cls() is DiffusionDecisionTrainer
training_args = SimpleNamespace(
pretraining=False,
sample_packing=True,
eval_sample_packing=False,
)
collator = builder.build_collator(
training_args,
padding=True,
pad_to_multiple_of=8192,
)
eval_collator = builder.build_collator(
training_args,
is_eval=True,
padding=True,
pad_to_multiple_of=8192,
)
for candidate in (collator, eval_collator):
assert isinstance(candidate, DecisionTrainingCollator)
assert candidate.spec is profile.diffusion
assert candidate.pad_token_id == 0
def test_config_accepts_versioned_sampled_slot_policy():
config = DiffusionDecisionConfig.model_validate(
{"latent": {"mode": "pinned", "num_slots": 8, "sample_num_slots": True}}
)
assert config.latent.sample_num_slots is True
assert config.latent.num_slots == 8
@pytest.mark.parametrize(
"latent,message",
[
(
{"mode": "none", "sample_num_slots": True},
"sample_num_slots requires a non-none",
),
(
{"mode": "pinned", "sample_num_slots": True},
"sample_num_slots requires num_slots",
),
],
)
def test_config_rejects_invalid_sampled_slot_bounds(latent, message):
with pytest.raises(ValidationError, match=message):
DiffusionDecisionConfig.model_validate({"latent": latent})
def test_post_lora_merge_ignores_manifests_for_non_decision_runs(tmp_path):
from axolotl.integrations.diffusion_decision.manifest import MANIFEST_FILENAME
from axolotl.integrations.diffusion_decision.plugin import DiffusionDecisionPlugin
adapter, merged = tmp_path / "adapter", tmp_path / "merged"
adapter.mkdir()
merged.mkdir()
(adapter / MANIFEST_FILENAME).write_text("{not a manifest")
DiffusionDecisionPlugin().post_lora_merge({}, str(adapter), str(merged))
assert not (merged / MANIFEST_FILENAME).exists()
def test_config_caps_questions_per_canvas_at_the_template_limit():
from axolotl.integrations.diffusion_decision.args import DiffusionDecisionConfig
from axolotl.integrations.diffusion_decision.vendored.djev_template import (
MAX_QUESTIONS,
)
DiffusionDecisionConfig.model_validate({"max_questions_per_canvas": MAX_QUESTIONS})
with pytest.raises(ValueError, match="max_questions_per_canvas"):
DiffusionDecisionConfig.model_validate(
{"max_questions_per_canvas": MAX_QUESTIONS + 1}
)
def test_collator_hook_defers_to_defaults_without_a_decision_block():
from axolotl.integrations.diffusion_decision.plugin import DiffusionDecisionPlugin
assert DiffusionDecisionPlugin().get_collator_cls_and_kwargs({}) is None
tests/integrations/diffusion_decision/test_preprocessing.py
+235
View File
@@ -0,0 +1,235 @@
"""Preprocessing uses independent upstream fixtures with actual tokenizer IDs."""
import random
import pytest
from axolotl.integrations.diffusion_decision.adapters import normalize_record
from axolotl.integrations.diffusion_decision.adapters.jsonl import normalize_jsonl
from axolotl.integrations.diffusion_decision.preprocessing import build_decision_canvas
from axolotl.integrations.diffusion_decision.records import OrdinalMetadata
from axolotl.integrations.diffusion_decision.slots import SlotInit
from tests.integrations.diffusion_decision.helpers import (
DJEV_FIXTURE,
RecordedTokenizer,
make_spec,
)
def inputs(family, index):
reference = DJEV_FIXTURE["tokenizers"][family]
case = reference["cases"][index]
labels = {}
for i, q in enumerate(case["parsed"]["questions"]):
count = len(q["labels"])
labels[q["id"]] = (
{"kind": "hard", "gold_idx": 0}
if i % 3 == 0
else {"kind": "dist", "probs": [1 / count] * count}
if i % 3 == 1
else {"kind": "set", "allowed_set": [0, count - 1]}
)
record = {
"source": "fixture",
"group": "synthetic",
"state": "observed",
"questions": {
q["id"]: {k: v for k, v in q.items() if k != "id"}
for q in case["schema"]["questions"]
},
"labels": labels,
}
return reference, case, record
@pytest.mark.parametrize("family", ["gemma", "dream"])
@pytest.mark.parametrize("index", [0, 9, 10, 19])
def test_exact_template_positions_noise_and_target_preservation(family, index):
reference, case, record = inputs(family, index)
kwargs = dict(
scaffold_ids=case["head"],
turn_close_id=106 if family == "gemma" else 151643,
pad_id=0 if family == "gemma" else 151643,
vocab_size=262144 if family == "gemma" else 152064,
seed=23,
)
first = build_decision_canvas(
RecordedTokenizer(reference["encodings"]), record, [1, 2], **kwargs
)
second = build_decision_canvas(
RecordedTokenizer(reference["encodings"]), record, [1, 2], steps=2, **kwargs
)
base, slots = case["template"]
expected = (
base + [kwargs["turn_close_id"]] + [kwargs["pad_id"]] * (128 - len(base) - 1)
)
rng = random.Random(23)
for slot in slots:
expected[slot["pos"]] = rng.randrange(kwargs["vocab_size"])
assert list(first.canvas_ids) == expected
assert first.canvas_ids == second.canvas_ids
assert first.template_length == len(base)
assert first.targets == tuple(record["labels"].values())
assert not any(first.pinned_mask)
assert all(
second.pinned_mask[i] == (i not in first.label_positions) for i in range(128)
)
assert all(first.semantic_mask) and not any(first.slot_mask)
def test_prevalidated_record_matches_validated_canvas():
reference, case, record = inputs("gemma", 0)
kwargs = dict(
scaffold_ids=case["head"],
turn_close_id=106,
pad_id=0,
vocab_size=262144,
seed=23,
)
tokenizer = RecordedTokenizer(reference["encodings"])
expected = build_decision_canvas(tokenizer, record, [1, 2], **kwargs)
actual = build_decision_canvas(
tokenizer,
normalize_jsonl(record),
[1, 2],
prevalidated_record=True,
**kwargs,
)
assert actual == expected
def test_absorbing_noise_and_invalid_soft_target():
reference, case, record = inputs("dream", 2)
kwargs = dict(
scaffold_ids=case["head"],
turn_close_id=151643,
pad_id=151643,
vocab_size=152064,
noise_kind="absorbing",
mask_token_id=151666,
)
result = build_decision_canvas(
RecordedTokenizer(reference["encodings"]), record, [1], **kwargs
)
assert all(result.canvas_ids[pos] == 151666 for pos in result.label_positions)
record["labels"]["1"]["probs"] = [0.1]
with pytest.raises(ValueError, match="one probability"):
build_decision_canvas(
RecordedTokenizer(reference["encodings"]), record, [1], **kwargs
)
@pytest.mark.parametrize("steps", [2, 3])
def test_generated_free_slots_remain_updateable_while_kstep_scaffold_is_pinned(steps):
reference, case, record = inputs("gemma", 0)
plan = SlotInit(
"free",
num_slots=2,
vocab_size=262144,
spec=make_spec(max_canvas=None, max_context=None),
).build(seed=7)
canvas = build_decision_canvas(
RecordedTokenizer(reference["encodings"]),
record,
[1, 2],
scaffold_ids=case["head"],
turn_close_id=106,
pad_id=0,
vocab_size=262144,
steps=steps,
slot_plan=plan,
)
slots = [index for index, value in enumerate(canvas.slot_mask) if value]
labels = set(canvas.label_positions)
assert len(slots) == 2
assert all(not canvas.pinned_mask[index] for index in slots)
assert all(not canvas.pinned_mask[index] for index in labels)
assert all(
canvas.pinned_mask[index]
for index in range(len(canvas.canvas_ids))
if index not in labels and index not in slots
)
@pytest.mark.parametrize("family", ["gemma", "dream"])
def test_canvas_builds_source_prompt_from_validated_record(family):
reference, case, record = inputs(family, 0)
record["state"] = {"observed": "20°"}
record["instructions"] = case["schema"]["instructions"]
class ChatTokenizer(RecordedTokenizer):
def apply_chat_template(self, messages, **kwargs):
assert messages == [
{"role": "system", "content": case["system_text"]},
{"role": "user", "content": '{"observed": "20°"}'},
]
assert kwargs == {
"tokenize": True,
"add_generation_prompt": True,
"enable_thinking": False,
}
return {"input_ids": [1, 2, 3]}
result = build_decision_canvas(
ChatTokenizer(reference["encodings"]),
record,
scaffold_ids=case["head"],
turn_close_id=106 if family == "gemma" else 151643,
pad_id=0 if family == "gemma" else 151643,
vocab_size=262144 if family == "gemma" else 152064,
)
assert result.prompt_ids == (1, 2, 3)
assert result.targets == tuple(record["labels"].values())
def test_soft_score_canvas_opt_in_ordinal_metadata(monkeypatch):
import axolotl.integrations.diffusion_decision.preprocessing as preprocessing
normalized = normalize_record(
"jsonl",
{
"id": "soft-score",
"source": "score_eval",
"group": "soft-score",
"state": {"evidence": "text"},
"questions": {
"q1": {
"type": "score",
"instructions": "",
"levels": ["low", "medium", "high"],
}
},
"labels": {"q1": {"kind": "dist", "probs": [0.2, 0.5, 0.3]}},
},
training=False,
)
monkeypatch.setattr(
preprocessing,
"resolve_template",
lambda *_args, **_kwargs: (
[9, 31, 11],
[{"pos": 1, "label_ids": [31, 32, 33]}],
),
)
kwargs = dict(
prompt_ids=(1,),
scaffold_ids=(),
turn_close_id=2,
pad_id=0,
vocab_size=64,
width=8,
)
default = build_decision_canvas(object(), normalized, **kwargs)
assert default.ordinal_metadata == ()
canvas = build_decision_canvas(
object(), normalized, include_ordinal_metadata=True, **kwargs
)
assert canvas.ordinal_metadata == (
OrdinalMetadata(
levels=("low", "medium", "high"),
source_ids=("0", "1", "2"),
candidate_ranks=(0, 1, 2),
),
)
tests/integrations/diffusion_decision/test_procedural.py
+100
View File
@@ -0,0 +1,100 @@
import json
import pytest
from axolotl.integrations.diffusion_decision.adapters.procedural import normalize
def test_soft_bool_preserved():
r = {
"id": "x:1",
"state": "s",
"questions": json.dumps({"q": {"type": "noul", "instructions": "x"}}),
"answers": json.dumps({"q": {"noul": 2 / 3}}),
}
x = normalize(r)
assert x["labels"]["q"]["kind"] == "dist"
assert x["labels"]["q"]["probs"] == pytest.approx([2 / 3, 1 / 3])
def test_choice_identity():
r = {
"id": "x:1",
"state": "s",
"questions": json.dumps(
{"q": {"type": "choice", "criteria": {"b": "b", "a": "a"}}}
),
"answers": json.dumps({"q": {"probabilities": {"b": 0, "a": 1}}}),
}
assert normalize(r)["labels"]["q"]["gold_idx"] == 1
def test_missing_noul_rejected():
r = {
"id": "x:1",
"state": "s",
"questions": '{"q":{"type":"noul"}}',
"answers": '{"q":{}}',
}
with pytest.raises(ValueError, match="missing noul"):
normalize(r)
def test_choice_description_retained():
r = {
"id": "x:1",
"state": "s",
"questions": '{"q":{"type":"choice","criteria":{"a":"alpha","b":"beta"}}}',
"answers": '{"q":{"probabilities":{"a":1,"b":0}}}',
}
assert normalize(r)["questions"]["q"]["options"] == [
{"name": "a", "description": "alpha"},
{"name": "b", "description": "beta"},
]
def test_noul_criteria_retained():
r = {
"id": "x:1",
"state": "s",
"questions": '{"q":{"type":"noul","criteria":{"true":"yes desc","false":"no desc"}}}',
"answers": '{"q":{"noul":1}}',
}
assert normalize(r)["questions"]["q"]["criteria"]["true"] == "yes desc"
def test_source_id_is_the_group_without_an_invented_task_family():
result = normalize(
{
"id": "arithmetic:train:7",
"state": "s",
"questions": '{"q":{"type":"noul"}}',
"answers": '{"q":{"noul":1}}',
}
)
assert result["group"] == "arithmetic:train:7"
assert "family" not in result
assert result["source_metadata"] == {
"task": "arithmetic",
"official_split": "train",
}
def test_score_legend_identity_overrides_dictionary_order():
row = {
"id": "task:train:1",
"state": "{}",
"questions": {"q": {"type": "score", "criteria": ["low", "high"]}},
"answers": {
"q": {
"legend": {"0": "high", "1": "low"},
"probabilities": {"0": 0.8, "1": 0.2},
}
},
}
result = normalize(row)
assert result["labels"]["q"]["probs"] == [0.2, 0.8]
assert result["questions"]["q"]["levels"] == ["low", "high"]
row["answers"]["q"]["legend"]["1"] = "high"
with pytest.raises(ValueError, match="exactly once"):
normalize(row)
tests/integrations/diffusion_decision/test_prompts.py
+75
View File
@@ -0,0 +1,75 @@
from typing import Any
import pytest
from axolotl.integrations.diffusion_decision import prompts
class RecordingTokenizer:
def __init__(self, output: Any):
self.output = output
self.messages: Any = None
self.kwargs: dict[str, Any] | None = None
def apply_chat_template(self, messages, **kwargs):
self.messages = messages
self.kwargs = kwargs
return self.output
def test_serializes_state_like_pinned_djev():
assert prompts.serialize_state("raw °") == "raw °"
assert prompts.serialize_state({"temperature": "20°"}) == '{"temperature": "20°"}'
with pytest.raises(ValueError, match="state is required"):
prompts.serialize_state(None)
def test_structured_content_and_dict_tokenizer_output():
tokenizer = RecordingTokenizer({"input_ids": [[4, 5, 6]]})
ids = prompts.decision_prompt_ids(
tokenizer,
"system",
{"value": "20°"},
thinking=True,
structured_content=True,
)
assert ids == (4, 5, 6)
assert tokenizer.messages == [
{"role": "system", "content": [{"type": "text", "text": "system"}]},
{
"role": "user",
"content": [{"type": "text", "text": '{"value": "20°"}'}],
},
]
assert tokenizer.kwargs == {
"tokenize": True,
"add_generation_prompt": True,
"enable_thinking": True,
}
def test_defaults_to_source_faithful_string_content():
tokenizer = RecordingTokenizer([7, 8])
assert prompts.decision_prompt_ids(tokenizer, "system", "raw state") == (7, 8)
assert tokenizer.messages == [
{"role": "system", "content": "system"},
{"role": "user", "content": "raw state"},
]
def test_rejects_multiple_prompts_or_non_token_output():
with pytest.raises(ValueError, match="more than one prompt"):
prompts.decision_prompt_ids(
RecordingTokenizer([[1], [2]]), "system", "state", structured_content=False
)
with pytest.raises(TypeError, match="token ids"):
prompts.decision_prompt_ids(
RecordingTokenizer(object()), "system", "state", structured_content=False
)
with pytest.raises(TypeError, match="integer token ids"):
prompts.decision_prompt_ids(
RecordingTokenizer([1.5]), "system", "state", structured_content=False
)
tests/integrations/diffusion_decision/test_public_decision_mix.py
+154
View File
@@ -0,0 +1,154 @@
import hashlib
import json
from importlib.util import module_from_spec, spec_from_file_location
from pathlib import Path
import pytest
from datasets import Dataset
SCRIPT = Path(__file__).parents[3] / "scripts/diffusion_lm/build_public_decision_mix.py"
SPEC = spec_from_file_location("build_public_decision_mix", SCRIPT)
assert SPEC and SPEC.loader
MODULE = module_from_spec(SPEC)
SPEC.loader.exec_module(MODULE)
def _raw(config: str, split: str, index: int) -> dict:
return {
"id": f"{config}:{split}:{index}",
"state": json.dumps({"config": config, "split": split, "index": index}),
"questions": json.dumps(
{"ok": {"type": "noul", "instructions": "", "criteria": {}}}
),
"answers": json.dumps({"ok": {"noul": 1.0}}),
}
def test_build_public_mix_is_deterministic_and_decontaminates_heldouts(
tmp_path, monkeypatch
):
def fake_load_dataset(dataset, name, split, revision):
assert (dataset, revision) == (MODULE.DATASET, MODULE.REVISION)
count = (
MODULE.TRAIN_PER_CONFIG + 1 if split == "train" else MODULE.DEV_PER_CONFIG
)
return Dataset.from_list([_raw(name, split, index) for index in range(count)])
monkeypatch.setattr(MODULE, "load_dataset", fake_load_dataset)
protected = tmp_path / "protected.jsonl"
colliding_id = min(
(f"arithmetic:train:{index}" for index in range(MODULE.TRAIN_PER_CONFIG + 1)),
key=lambda record_id: hashlib.sha256(
f"{MODULE.SEED}\0arithmetic\0{record_id}".encode()
).hexdigest(),
)
colliding_index = int(colliding_id.rsplit(":", 1)[1])
protected.write_text(
"\n".join(
json.dumps(row)
for row in (
{"source": "eval_a", "family": "heldout", "state": {"x": "j"}},
{
"source": "eval_a",
"family": "heldout",
"state": {
"config": "arithmetic",
"split": "train",
"index": colliding_index,
},
},
)
)
+ "\n"
)
first = MODULE.build(tmp_path / "first", [protected])
second = MODULE.build(tmp_path / "second", [protected])
assert first["files"] == second["files"]
assert first["protected_decontamination"]["status"] == "performed"
assert first["protected_decontamination"]["protected_rows_by_source"] == {
"eval_a": 2
}
assert (
first["files"]["train"]["lines"]
== len(MODULE.CONFIGS) * MODULE.TRAIN_PER_CONFIG
)
assert first["files"]["dev"]["lines"] == len(MODULE.CONFIGS) * MODULE.DEV_PER_CONFIG
assert (
first["files"]["train"]["sha256"]
== hashlib.sha256((tmp_path / "first" / "train.jsonl").read_bytes()).hexdigest()
)
assert '"source":"eval_a"' not in (tmp_path / "first" / "train.jsonl").read_text()
train_ids = {
json.loads(line)["id"]
for line in (tmp_path / "first" / "train.jsonl").read_text().splitlines()
}
assert colliding_id not in train_ids
assert len(train_ids) == first["files"]["train"]["lines"]
def test_build_public_mix_is_usable_without_protected_rows(tmp_path, monkeypatch):
def fake_load_dataset(dataset, name, split, revision):
count = MODULE.TRAIN_PER_CONFIG if split == "train" else MODULE.DEV_PER_CONFIG
return Dataset.from_list([_raw(name, split, index) for index in range(count)])
monkeypatch.setattr(MODULE, "load_dataset", fake_load_dataset)
contract = MODULE.build(tmp_path / "out")
assert contract["protected_decontamination"]["status"] == "not_requested"
def _wide_choice(count):
options = {str(index): f"option {index}" for index in range(count)}
return {
"id": "arithmetic:train:0",
"state": json.dumps({"x": 1}),
"questions": json.dumps(
{"choice": {"type": "choice", "instructions": "", "criteria": options}}
),
"answers": json.dumps(
{
"choice": {
"probabilities": {
name: 0.5 if name in {"0", "1"} else 0.0 for name in options
}
}
}
),
}
def test_normalize_uses_recipe_codebook(monkeypatch):
rows = [_wide_choice(26)]
monkeypatch.setattr(
MODULE, "load_dataset", lambda *args, **kwargs: Dataset.from_list(rows)
)
record = MODULE._normalize("arithmetic", "train")[0]
assert record["labels"]["choice"]["probs"][0] == 0.5
rows[:] = [_wide_choice(27)]
with pytest.raises(ValueError, match="at most 26 alternatives"):
MODULE._normalize("arithmetic", "train")
def test_source_group_overlap_is_rejected():
record = {"source": "public_procedural.arithmetic", "group": "origin-1"}
with pytest.raises(ValueError, match="source/group"):
MODULE._assert_source_group_isolation({"train": [record], "dev": [record]})
def test_build_public_mix_rejects_duplicate_source_ids(tmp_path, monkeypatch):
def fake_load_dataset(dataset, name, split, revision):
count = MODULE.TRAIN_PER_CONFIG if split == "train" else MODULE.DEV_PER_CONFIG
rows = [_raw(name, split, index) for index in range(count)]
if name == "arithmetic" and split == "train":
rows[-1]["id"] = rows[0]["id"]
return Dataset.from_list(rows)
monkeypatch.setattr(MODULE, "load_dataset", fake_load_dataset)
protected = tmp_path / "protected.jsonl"
protected.write_text(
json.dumps({"source": "eval_a", "family": "heldout", "state": {"x": "j"}})
+ "\n"
)
with pytest.raises(ValueError, match="duplicate public source ids"):
MODULE.build(tmp_path / "out", [protected])
tests/integrations/diffusion_decision/test_reader_validation.py
+82
View File
@@ -0,0 +1,82 @@
"""Reject malformed canvas metadata before invoking a model or serving API."""
from dataclasses import replace
import pytest
import torch
from axolotl.integrations.diffusion_decision.readers.base import (
restricted_probabilities,
validate_canvas,
)
from tests.integrations.diffusion_decision.helpers import (
make_canvas,
)
def canvas():
return make_canvas(
(1, 2),
(3, 4, 5, 0),
(1,),
allowed=(6, 7),
question_ids=("answer",),
template_length=2,
)
@pytest.mark.parametrize(
"changes,message",
[
({"prompt_ids": ()}, "requires a prompt"),
({"canvas_ids": (3, -1, 5, 0)}, "nonnegative"),
({"allowed_ids": ((True, 7),)}, "nonnegative"),
({"template_length": 4}, "template_length"),
({"label_positions": (2,)}, "outside"),
({"slot_mask": (False, True, False, False)}, "latent slots"),
({"question_ids": ("",)}, "distinct nonempty"),
],
)
def test_invalid_canvas_rejected(changes, message):
with pytest.raises(ValueError, match=message):
validate_canvas(replace(canvas(), **changes))
def test_duplicate_question_ids_rejected():
with pytest.raises(ValueError, match="distinct nonempty"):
validate_canvas(replace(canvas(), question_ids=("answer", "answer")))
def test_full_width_pad_tokens_remain_semantically_visible():
assert validate_canvas(canvas()) == 4
def test_candidate_mask_does_not_extract_tensor_scalars(monkeypatch):
logits = torch.log_softmax(torch.tensor([[1.0, 2.0, 3.0], [3.0, 2.0, 1.0]]), -1)
def fail(*_args, **_kwargs):
raise AssertionError("candidate mask extracted a tensor scalar")
with monkeypatch.context() as patch:
for name in ("item", "__bool__", "__int__"):
patch.setattr(torch.Tensor, name, fail)
_, mask, probabilities = restricted_probabilities(logits, ((0, 2), (1,)))
assert mask.tolist() == [[True, True], [True, False]]
torch.testing.assert_close(
probabilities[0], torch.softmax(torch.tensor([1.0, 3.0]), 0)
)
torch.testing.assert_close(probabilities[1], torch.tensor([1.0, 0.0]))
@pytest.mark.parametrize("token", [-1, 3])
def test_candidate_token_outside_vocabulary_rejected(token):
with pytest.raises(ValueError, match="outside the model vocabulary"):
restricted_probabilities(torch.zeros(1, 3), ((token,),))
def test_validate_canvas_rejects_pinned_label_positions():
pinned = list(canvas().pinned_mask)
pinned[1] = True
with pytest.raises(ValueError, match="pinned tokens cannot be label positions"):
validate_canvas(replace(canvas(), pinned_mask=tuple(pinned)))
tests/integrations/diffusion_decision/test_slot_canvases.py
+131
View File
@@ -0,0 +1,131 @@
"""Slot plans affect only declared decision-canvas regions."""
import pytest
from axolotl.integrations.diffusion_decision.template import SchemaError
from axolotl.model_support import (
DiffusionNoise,
)
from tests.integrations.diffusion_decision.helpers import (
build_canvas,
free_plan_seed,
make_slot_plan,
)
def _plan(mode, *, count=2, ids=(), noise=DiffusionNoise.UNIFORM):
return make_slot_plan(
mode, count=count, ids=ids, noise=noise, seed=free_plan_seed(mode, noise)
)
def test_none_slot_plan_keeps_existing_canvas_bytes_even_with_explicit_tags():
baseline = build_canvas()
none = build_canvas(_plan("none", count=0))
assert none == baseline
@pytest.mark.parametrize(
("mode", "ids", "placement", "expected_slots"),
[
("pad", (), "after_turn", (0, 0)),
("pinned", (7,), "thought", (7, 7)),
("pinned", (7, 8), "thought", (7, 8)),
("learned", (7, 8), "thought", (7, 8)),
("prompt", (7, 8), "prompt", (7, 8)),
],
)
def test_slot_modes_preserve_labels_and_mark_only_canvas_slots(
mode, ids, placement, expected_slots
):
plan = _plan(mode, ids=ids)
canvas = build_canvas(plan)
baseline = build_canvas()
assert (
canvas.canvas_ids[canvas.label_positions[0]]
== baseline.canvas_ids[baseline.label_positions[0]]
)
assert canvas.allowed_ids == baseline.allowed_ids
assert canvas.targets == baseline.targets
assert len(canvas.canvas_ids) == len(baseline.canvas_ids) == 128
if placement == "thought":
assert canvas.canvas_ids[:5] == (70, 71, *expected_slots, 72)
assert [index for index, value in enumerate(canvas.slot_mask) if value] == [
2,
3,
]
assert canvas.pinned_mask[2:4] == (True, True)
elif placement == "after_turn":
assert canvas.canvas_ids[canvas.template_length] == 106
assert (
canvas.canvas_ids[canvas.template_length + 1 : canvas.template_length + 3]
== expected_slots
)
assert [index for index, value in enumerate(canvas.slot_mask) if value] == [
canvas.template_length + 1,
canvas.template_length + 2,
]
assert all(
canvas.pinned_mask[index]
for index, value in enumerate(canvas.slot_mask)
if value
)
else:
assert canvas.prompt_ids == (*expected_slots, 99)
assert canvas.prompt_slot_mask == (True, True, False)
assert not any(canvas.slot_mask)
assert canvas.canvas_ids == baseline.canvas_ids
if placement != "prompt":
assert canvas.prompt_slot_mask == ()
assert all(canvas.semantic_mask)
assert all(not canvas.slot_mask[position] for position in canvas.label_positions)
def test_mask_and_free_slots_follow_noise_specific_initialization_and_pin_policy():
masked = build_canvas(
_plan("mask", noise=DiffusionNoise.ABSORBING), noise_kind="absorbing"
)
free_plan = _plan("free")
free = build_canvas(free_plan)
assert masked.canvas_ids[:5] == (70, 71, 9, 9, 72)
assert masked.pinned_mask[2:4] == (True, True)
assert masked.canvas_ids[masked.label_positions[0]] == 9
assert free.slot_mask[2:4] == (True, True)
assert free.pinned_mask[2:4] == (False, False)
assert free.canvas_ids[2:4] == free_plan.ids
def test_fixed_slots_stay_pinned_at_one_read_and_non_slot_behavior_stays_source_faithful():
pinned = build_canvas(_plan("pinned", ids=(7,)), steps=1)
repeated = build_canvas(_plan("pinned", ids=(7,)), steps=2)
assert pinned.pinned_mask[2:4] == (True, True)
assert pinned.pinned_mask[pinned.label_positions[0]] is False
assert repeated.pinned_mask[pinned.label_positions[0]] is False
assert all(
repeated.pinned_mask[index] == (index not in repeated.label_positions)
for index in range(128)
)
def test_distinct_pinned_slots_match_learned_canvas_coordinates_without_trainable_rows():
pinned_plan = _plan("pinned", ids=(7, 8))
learned_plan = _plan("learned", ids=(7, 8))
pinned = build_canvas(pinned_plan, steps=2)
learned = build_canvas(learned_plan, steps=2)
assert pinned.canvas_ids == learned.canvas_ids
assert pinned.label_positions == learned.label_positions
assert pinned.slot_mask == learned.slot_mask
assert pinned.pinned_mask == learned.pinned_mask
assert pinned_plan.trainable_token_ids == ()
assert learned_plan.trainable_token_ids == (7, 8)
def test_slot_budget_overflow_raises_schema_error_without_truncation():
with pytest.raises(SchemaError, match="canvas"):
build_canvas(_plan("learned", ids=(7, 8)), width=7)
tests/integrations/diffusion_decision/test_slot_collation.py
+174
View File
@@ -0,0 +1,174 @@
"""Fixed decision slots retain their canvas and collation invariants."""
from __future__ import annotations
import pytest
import torch
from axolotl.integrations.diffusion_decision.training_collator import (
DecisionTrainingCollator,
)
from axolotl.model_support import (
DiffusionLayout,
DiffusionNoise,
)
from tests.integrations.diffusion_decision.helpers import (
build_canvas,
make_record,
make_rows,
make_slot_plan,
make_spec,
slot_token_ids,
)
def _canvas(identifier: str, mode: str, noise: DiffusionNoise):
return build_canvas(
make_slot_plan(mode, ids=slot_token_ids(mode), noise=noise),
record=make_record(identifier, state=f"state-{identifier}"),
prompt_ids=(90 + int(identifier),),
noise=noise,
)
def _rows(mode: str, noise: DiffusionNoise):
canvases = tuple(_canvas(str(index), mode, noise) for index in range(2))
return canvases, make_rows(canvases)
@pytest.mark.parametrize(
("mode", "noise", "expected_ids"),
[
("pinned", DiffusionNoise.UNIFORM, (7, 7)),
("learned", DiffusionNoise.UNIFORM, (7, 8)),
("mask", DiffusionNoise.ABSORBING, (9, 9)),
("pad", DiffusionNoise.UNIFORM, (0, 0)),
("prompt", DiffusionNoise.UNIFORM, (7, 8)),
],
)
@pytest.mark.parametrize(
"layout", [DiffusionLayout.FULL_SEQUENCE, DiffusionLayout.ENCODER_CANVAS]
)
def test_fixed_slot_collation_keeps_slots_pinned_and_out_of_canvas_loss(
mode: str,
noise: DiffusionNoise,
expected_ids: tuple[int, ...],
layout: DiffusionLayout,
):
canvases, rows = _rows(mode, noise)
batch = DecisionTrainingCollator(make_spec(layout=layout, noise=noise))(rows)
for canvas in canvases:
slot_positions = tuple(
index for index, value in enumerate(canvas.slot_mask) if value
)
assert set(slot_positions).isdisjoint(canvas.label_positions)
if mode == "prompt":
assert canvas.prompt_ids[: len(expected_ids)] == expected_ids
assert canvas.prompt_slot_mask == (True, True, False)
assert not any(canvas.slot_mask)
else:
assert (
tuple(canvas.canvas_ids[index] for index in slot_positions)
== expected_ids
)
assert all(canvas.pinned_mask[index] for index in slot_positions)
if layout is DiffusionLayout.FULL_SEQUENCE:
_assert_full_sequence_contract(batch, canvases, mode, expected_ids)
else:
_assert_encoder_canvas_contract(
batch["diffusion_batch"], canvases, mode, expected_ids
)
def _assert_full_sequence_contract(batch, canvases, mode, expected_ids):
input_ids = batch["input_ids"][0]
pinned = batch["canvas_input_pinned_mask"][0]
loss = batch["canvas_loss_mask"][0]
corruptible = batch["canvas_corruptible_mask"][0]
update = batch["canvas_update_mask"][0]
offset = 0
for document, canvas in enumerate(canvases):
prompt_length = len(canvas.prompt_ids)
canvas_start = offset + prompt_length
slot_positions = tuple(
index for index, value in enumerate(canvas.slot_mask) if value
)
fixed_positions = (
tuple(range(offset, offset + len(expected_ids)))
if mode == "prompt"
else tuple(canvas_start + index for index in slot_positions)
)
assert (
tuple(input_ids[position].item() for position in fixed_positions)
== expected_ids
)
assert torch.all(pinned[list(fixed_positions)])
assert not torch.any(loss[list(fixed_positions)])
assert not torch.any(corruptible[list(fixed_positions)])
assert not torch.any(update[list(fixed_positions)])
assert batch["document_ids"][
0, offset : offset + prompt_length + len(canvas.canvas_ids)
].tolist() == [document] * (prompt_length + len(canvas.canvas_ids))
assert batch["position_ids"][
0, offset : offset + prompt_length + len(canvas.canvas_ids)
].tolist() == list(range(prompt_length + len(canvas.canvas_ids)))
if mode == "prompt":
assert max(fixed_positions) < canvas_start
assert canvas.prompt_slot_mask == (True, True, False)
offset += prompt_length + len(canvas.canvas_ids)
def _assert_encoder_canvas_contract(diffusion_batch, canvases, mode, expected_ids):
for row, canvas in enumerate(canvases):
slot_positions = tuple(
index for index, value in enumerate(canvas.slot_mask) if value
)
if mode == "prompt":
fixed_positions = tuple(range(len(expected_ids)))
assert canvas.prompt_slot_mask == (True, True, False)
assert (
tuple(
diffusion_batch.encoder_input_ids[
row, list(fixed_positions)
].tolist()
)
== expected_ids
)
assert torch.all(
diffusion_batch.encoder_validity[row, list(fixed_positions)]
)
assert not torch.any(
diffusion_batch.encoder_ar_valid_mask[row, list(fixed_positions)]
)
assert diffusion_batch.encoder_ar_valid_mask[row, len(expected_ids)]
else:
fixed_positions = slot_positions
assert (
tuple(
diffusion_batch.canvas_clean_ids[
row, list(fixed_positions)
].tolist()
)
== expected_ids
)
assert torch.all(
diffusion_batch.canvas_input_pinned_mask[row, list(slot_positions)]
)
assert not torch.any(
diffusion_batch.canvas_loss_mask[row, list(slot_positions)]
)
assert not torch.any(
diffusion_batch.canvas_corruptible_mask[row, list(slot_positions)]
)
assert not torch.any(
diffusion_batch.canvas_update_mask[row, list(slot_positions)]
)
assert not torch.any(
diffusion_batch.canvas_sc_eligible_mask[row, list(slot_positions)]
)
assert torch.all(
diffusion_batch.canvas_read_only_mask[row, list(slot_positions)]
)
tests/integrations/diffusion_decision/test_slot_runtime.py
+188
View File
@@ -0,0 +1,188 @@
"""Learned decision slots wire only audited embedding rows into PEFT."""
import pytest
import torch
from peft import LoraConfig, TaskType, get_peft_model
from transformers import LlamaConfig, LlamaForCausalLM
from axolotl.integrations.diffusion_decision.args import DiffusionDecisionConfig
from axolotl.integrations.diffusion_decision.plugin import DiffusionDecisionPlugin
from axolotl.integrations.diffusion_decision.slot_runtime import (
merge_trainable_slot_indices,
resolve_trainable_slot_runtime,
)
from axolotl.model_support import (
DiffusionNoise,
FirstPositionAlignment,
ReductionScope,
TimeWeighting,
)
from axolotl.utils.dict import DictDefault
from tests.integrations.diffusion_decision.helpers import (
make_spec,
)
def _spec():
return make_spec(
noise=DiffusionNoise.ABSORBING,
first_position_alignment=FirstPositionAlignment.REQUIRES_PREDECESSOR,
max_canvas=None,
max_context=64,
default_time_weighting=TimeWeighting.INV_T,
reduction_scope=ReductionScope.GLOBAL_WINDOW,
)
def _model():
config = LlamaConfig(
vocab_size=131072,
hidden_size=16,
intermediate_size=32,
num_hidden_layers=1,
num_attention_heads=2,
num_key_value_heads=2,
pad_token_id=None,
bos_token_id=1,
eos_token_id=11,
)
config.mask_token_id = 100
return LlamaForCausalLM(config)
def _cfg(*, mode="learned", ids=(936, 937), existing=None):
return DictDefault(
{
"model_config_type": "nemotron_labs_diffusion",
"diffusion_lm": {"from_causal_lm": False},
"diffusion_decision": {
"latent": {"mode": mode, "num_slots": len(ids), "token_ids": ids}
},
"peft_trainable_token_indices": existing,
}
)
def _decision(cfg):
return DiffusionDecisionConfig.model_validate(cfg.diffusion_decision)
def test_pre_lora_merges_learned_rows_before_real_peft_embedding_wrap():
torch.manual_seed(0)
cfg = _cfg(existing=[900])
model = _model()
base_weight = model.get_input_embeddings().weight
outside_before = base_weight.detach().clone()
DiffusionDecisionPlugin().pre_lora_load(cfg, model)
assert cfg.peft_trainable_token_indices == [900, 936, 937]
wrapped = get_peft_model(
model,
LoraConfig(
task_type=TaskType.CAUSAL_LM,
target_modules=["q_proj"],
r=2,
lora_alpha=2,
trainable_token_indices=cfg.peft_trainable_token_indices,
),
)
optimizer = torch.optim.SGD(
[parameter for parameter in wrapped.parameters() if parameter.requires_grad],
lr=0.1,
)
loss = wrapped(
input_ids=torch.tensor([[936, 937, 5]]),
labels=torch.tensor([[-100, -100, 5]]),
).loss
loss.backward()
token_deltas = [
parameter
for name, parameter in wrapped.named_parameters()
if "trainable_tokens_delta" in name
]
assert token_deltas
assert any(
parameter.grad is not None and torch.count_nonzero(parameter.grad) > 0
for parameter in token_deltas
)
optimizer.step()
assert base_weight.requires_grad is False
assert torch.equal(base_weight.detach(), outside_before)
def test_prompt_rows_use_actual_embedding_path_and_preserve_mapping_indices():
cfg = _cfg(mode="prompt", existing={"model.embed_tokens": [900]})
model = _model()
runtime = resolve_trainable_slot_runtime(cfg, model, _decision(cfg), spec=_spec())
assert runtime is not None
assert runtime.embedding_module_path == "model.embed_tokens"
assert runtime.plan.trainable_token_ids == (936, 937)
merge_trainable_slot_indices(cfg, runtime)
assert cfg.peft_trainable_token_indices == {"model.embed_tokens": [900, 936, 937]}
def test_nonlearned_modes_do_not_change_peft_embedding_rows():
cfg = _cfg(mode="pinned", ids=(936,), existing=[900])
model = _model()
DiffusionDecisionPlugin().pre_lora_load(cfg, model)
assert cfg.peft_trainable_token_indices == [900]
def test_distinct_pinned_rows_do_not_change_peft_embedding_rows():
cfg = _cfg(mode="pinned", ids=(936, 937), existing=[900])
model = _model()
DiffusionDecisionPlugin().pre_lora_load(cfg, model)
assert cfg.peft_trainable_token_indices == [900]
@pytest.mark.parametrize("token_id", (1, 11, 100))
def test_active_model_control_ids_are_rejected_for_trainable_slots(token_id):
cfg = _cfg(ids=(token_id, 936))
with pytest.raises(ValueError, match="active pad, bos, eos, or mask"):
resolve_trainable_slot_runtime(cfg, _model(), _decision(cfg), spec=_spec())
def test_mapping_without_actual_embedding_path_fails_explicitly():
cfg = _cfg(existing={"other.embedding": [900]})
runtime = resolve_trainable_slot_runtime(
cfg, _model(), _decision(cfg), spec=_spec()
)
assert runtime is not None
with pytest.raises(ValueError, match="actual input embedding path"):
merge_trainable_slot_indices(cfg, runtime)
def test_runtime_uses_embedding_vocab_not_model_config_vocab():
model = _model()
model.config.vocab_size = 4
cfg = _cfg(ids=(936, 937))
runtime = resolve_trainable_slot_runtime(cfg, model, _decision(cfg), spec=_spec())
assert runtime is not None
assert runtime.plan.ids == (936, 937)
def test_multiple_eos_ids_are_controls_without_requiring_a_pad_token():
model = _model()
model.config.eos_token_id = [1, 106]
runtime = resolve_trainable_slot_runtime(
_cfg(ids=(936, 937)), model, _decision(_cfg(ids=(936, 937))), spec=_spec()
)
assert runtime is not None
assert runtime.plan.ids == (936, 937)
with pytest.raises(ValueError, match="active pad, bos, eos, or mask"):
cfg = _cfg(ids=(106, 936))
resolve_trainable_slot_runtime(cfg, model, _decision(cfg), spec=_spec())
tests/integrations/diffusion_decision/test_slot_sampling.py
+208
View File
@@ -0,0 +1,208 @@
"""Pure maximum-canvas projection tests for sampled decision slot counts."""
from __future__ import annotations
import random
import pytest
from axolotl.integrations.diffusion_decision.slot_sampling import (
UNIFORM_INCLUSIVE_V1,
DecisionDraw,
project_max_canvas,
project_slot_plan,
sample_slot_count,
)
from axolotl.integrations.diffusion_decision.slots import SlotPlan
from axolotl.model_support import (
DiffusionNoise,
)
from tests.integrations.diffusion_decision.helpers import (
build_canvas,
free_plan_seed,
make_slot_plan,
slot_token_ids,
)
def _maximum_plan(mode: str, noise: DiffusionNoise) -> SlotPlan:
return make_slot_plan(
mode,
count=3,
ids=slot_token_ids(mode, 3),
noise=noise,
seed=free_plan_seed(mode, noise),
)
def _canvas(plan: SlotPlan, *, noise: DiffusionNoise, width: int = 128):
return build_canvas(plan, noise=noise, width=width, steps=2)
@pytest.mark.parametrize(
("mode", "noise"),
[
("pad", DiffusionNoise.UNIFORM),
("pinned", DiffusionNoise.UNIFORM),
("learned", DiffusionNoise.UNIFORM),
("prompt", DiffusionNoise.UNIFORM),
("mask", DiffusionNoise.ABSORBING),
("free", DiffusionNoise.UNIFORM),
],
)
@pytest.mark.parametrize("count", [0, 1, 3])
def test_projection_matches_direct_builder_for_all_supported_placements(
mode: str, noise: DiffusionNoise, count: int
):
maximum_plan = _maximum_plan(mode, noise)
maximum_canvas = _canvas(maximum_plan, noise=noise)
projected_plan = project_slot_plan(maximum_plan, count)
projected_canvas = project_max_canvas(
maximum_canvas,
maximum_plan,
count,
pad_token_id=0,
padding_pinned=True,
)
direct = _canvas(projected_plan, noise=noise)
assert projected_canvas == direct
assert projected_canvas.targets == maximum_canvas.targets
assert projected_canvas.allowed_ids == maximum_canvas.allowed_ids
assert projected_canvas.question_ids == maximum_canvas.question_ids
assert all(
not projected_canvas.slot_mask[position]
for position in projected_canvas.label_positions
)
assert all(not value for value in projected_plan.loss_mask)
if count == len(maximum_plan.ids):
assert projected_plan is maximum_plan
assert projected_canvas is maximum_canvas
def test_projection_preserves_free_slot_update_mask_and_typed_metadata():
maximum_plan = _maximum_plan("free", DiffusionNoise.UNIFORM)
maximum_canvas = _canvas(maximum_plan, noise=DiffusionNoise.UNIFORM)
plan = project_slot_plan(maximum_plan, 1)
canvas = project_max_canvas(
maximum_canvas,
maximum_plan,
1,
pad_token_id=0,
padding_pinned=True,
)
assert plan.update_mask == (True,)
assert plan.pinned_mask == (False,)
assert canvas.ordinal_metadata == maximum_canvas.ordinal_metadata
slot_position = canvas.slot_mask.index(True)
assert canvas.pinned_mask[slot_position] is False
assert canvas.canvas_ids[slot_position] == maximum_plan.ids[0]
@pytest.mark.parametrize(
"kwargs",
[
{"seed": True, "epoch": 0, "global_draw_ordinal": 0, "max_slots": 3},
{"seed": 1, "epoch": -1, "global_draw_ordinal": 0, "max_slots": 3},
{
"seed": 1,
"epoch": 0,
"global_draw_ordinal": 0,
"max_slots": 3,
"min_slots": 4,
},
{
"seed": 1,
"epoch": 0,
"global_draw_ordinal": 0,
"max_slots": 3,
"policy": "other",
},
{"seed": 1, "epoch": 0, "global_draw_ordinal": 0, "max_slots": 1 << 256},
],
)
def test_slot_count_sampler_rejects_invalid_contracts(kwargs):
with pytest.raises(ValueError):
sample_slot_count(**kwargs)
@pytest.mark.parametrize(
"kwargs",
[
{"index": -1, "epoch": 0, "global_draw_ordinal": 0},
{"index": 0, "epoch": True, "global_draw_ordinal": 0},
{"index": 0, "epoch": 0, "global_draw_ordinal": -1},
],
)
def test_decision_draw_rejects_invalid_stable_identity(kwargs):
with pytest.raises(ValueError):
DecisionDraw(**kwargs)
def test_slot_count_sampler_is_stateless_replayable_and_uniform(monkeypatch):
def forbidden_random(*_args, **_kwargs):
raise AssertionError("global RNG")
monkeypatch.setattr(random, "Random", forbidden_random)
first = [
sample_slot_count(seed=41, epoch=2, global_draw_ordinal=index, max_slots=4)
for index in range(10_000)
]
second = [
sample_slot_count(seed=41, epoch=2, global_draw_ordinal=index, max_slots=4)
for index in range(10_000)
]
changed_epoch = sample_slot_count(
seed=41, epoch=3, global_draw_ordinal=0, max_slots=4
)
counts = [first.count(value) for value in range(5)]
assert first == second
assert 0 <= changed_epoch <= 4
assert all(abs(value - 2_000) < 180 for value in counts)
assert UNIFORM_INCLUSIVE_V1 == "uniform_inclusive_v1"
def test_projection_rejects_invalid_count_and_malformed_maximum_canvas():
plan = _maximum_plan("learned", DiffusionNoise.UNIFORM)
canvas = _canvas(plan, noise=DiffusionNoise.UNIFORM)
with pytest.raises(ValueError, match="exceeds"):
project_slot_plan(plan, 4)
with pytest.raises(ValueError, match="exceeds"):
project_max_canvas(canvas, plan, 4, pad_token_id=0, padding_pinned=True)
malformed = SlotPlan(
ids=plan.ids,
placement=plan.placement,
pinned_mask=plan.pinned_mask,
update_mask=plan.update_mask,
loss_mask=(True, False, False),
trainable_token_ids=plan.trainable_token_ids,
)
with pytest.raises(ValueError, match="cannot carry loss"):
project_max_canvas(canvas, malformed, 1, pad_token_id=0, padding_pinned=True)
@pytest.mark.parametrize("mode", ["learned", "pad"])
def test_projection_handles_exact_fit_maximum_canvases(mode: str):
noise = DiffusionNoise.UNIFORM
plan = _maximum_plan(mode, noise)
padded = _canvas(plan, noise=noise)
width = padded.template_length + 1 + (len(plan.ids) if mode == "pad" else 0)
maximum = _canvas(plan, noise=noise, width=width)
assert maximum.slot_mask[-1] or maximum.canvas_ids[-1] != 0
projected = project_max_canvas(
maximum, plan, 1, pad_token_id=0, padding_pinned=True
)
direct = _canvas(project_slot_plan(plan, 1), noise=noise, width=width)
assert projected == direct
assert projected.canvas_ids[-2:] == (0, 0)
assert projected.slot_mask[-2:] == (False, False)
assert projected.pinned_mask[-2:] == (True, True)
tests/integrations/diffusion_decision/test_slots.py
+110
View File
@@ -0,0 +1,110 @@
"""Slot initialization contracts stay model-agnostic and pure."""
import random
import pytest
from axolotl.integrations.diffusion_decision.slots import SlotInit
from axolotl.model_support import (
DiffusionNoise,
)
from tests.integrations.diffusion_decision.helpers import (
make_slot_init,
make_spec,
)
def _init(mode, *, count=3, ids=(), noise=DiffusionNoise.UNIFORM, **kwargs):
return make_slot_init(
mode,
count=count,
ids=ids,
vocab_size=32,
spec=make_spec(noise=noise, self_conditioning=noise is DiffusionNoise.UNIFORM),
**kwargs,
)
def test_none_has_no_ids_and_no_slot_masks():
slot = SlotInit()
assert slot.ids == ()
assert slot.build().ids == ()
assert slot.build().placement == "none"
@pytest.mark.parametrize(
("mode", "ids", "placement", "expected", "trainable"),
[
("pad", (), "after_turn", (0, 0, 0), ()),
("pinned", (7,), "thought", (7, 7, 7), ()),
("pinned", (7, 8, 10), "thought", (7, 8, 10), ()),
("learned", (7, 8, 9), "thought", (7, 8, 9), (7, 8, 9)),
("prompt", (7, 8, 9), "prompt", (7, 8, 9), (7, 8, 9)),
],
)
def test_fixed_slot_modes_expose_placement_pinning_and_no_loss(
mode, ids, placement, expected, trainable
):
plan = _init(mode, ids=ids).build()
assert plan.ids == expected
assert plan.placement == placement
assert plan.pinned_mask == (True, True, True)
assert plan.update_mask == (False, False, False)
assert plan.loss_mask == (False, False, False)
assert plan.trainable_token_ids == trainable
def test_mask_is_absorbing_only_and_never_updates_or_contributes_loss():
plan = _init("mask", noise=DiffusionNoise.ABSORBING).build()
assert plan.ids == (9, 9, 9)
assert plan.placement == "thought"
assert plan.pinned_mask == (True, True, True)
assert plan.update_mask == (False, False, False)
assert plan.loss_mask == (False, False, False)
with pytest.raises(ValueError, match="absorbing"):
_init("mask").build()
def test_free_uniform_is_seeded_and_free_absorbing_starts_masked():
uniform = _init("free").build(seed=17)
generator = random.Random(17)
assert uniform.ids == tuple(generator.randrange(32) for _ in range(3))
assert uniform.placement == "thought"
assert uniform.pinned_mask == (False, False, False)
assert uniform.update_mask == (True, True, True)
assert uniform.loss_mask == (False, False, False)
assert _init("free").build(seed=17) == uniform
with pytest.raises(ValueError, match="explicit integer seed"):
_init("free").build()
with pytest.raises(ValueError, match="either seed or generator"):
_init("free").build(seed=17, generator=random.Random(17))
absorbing = _init("free", noise=DiffusionNoise.ABSORBING).build()
assert absorbing.ids == (9, 9, 9)
assert absorbing.pinned_mask == (False, False, False)
assert absorbing.update_mask == (True, True, True)
@pytest.mark.parametrize(
("slot", "message"),
[
(SlotInit("none", token_ids=(1,)), "mode=none"),
(_init("pad", ids=(1,)), "does not accept"),
(_init("pinned", ids=(1, 2)), "exactly 1 token_id or exactly num_slots"),
(_init("pinned", ids=(1, 1, 2)), "distinct token_ids"),
(_init("learned", ids=(1, 1, 2)), "distinct"),
(_init("prompt", ids=(1, 2)), "exactly 3"),
(_init("free", ids=(1,)), "does not accept"),
(_init("mask", ids=(1,), noise=DiffusionNoise.ABSORBING), "does not accept"),
(_init("pinned", ids=(32,)), "within the model vocabulary"),
(_init("pinned", count=0, ids=(1,)), "positive"),
],
)
def test_slot_modes_reject_invalid_counts_ids_and_vocab(slot, message):
with pytest.raises(ValueError, match=message):
slot.build()
tests/integrations/diffusion_decision/test_trainer.py
+2128
View File
File diff suppressed because it is too large Load Diff
tests/integrations/diffusion_decision/test_training_collator.py
+228
View File
@@ -0,0 +1,228 @@
from __future__ import annotations
from types import SimpleNamespace
import pytest
import torch
from torch.utils.data import DataLoader, SequentialSampler
from axolotl.integrations.diffusion_decision.datasets import DecisionDataset
from axolotl.integrations.diffusion_decision.records import DecisionCanvas
from axolotl.integrations.diffusion_decision.slot_sampling import DecisionDraw
from axolotl.integrations.diffusion_decision.training_collator import (
DecisionTrainingCollator,
)
from axolotl.model_support.diffusion import (
DiffusionLayout,
FirstPositionAlignment,
)
from tests.integrations.diffusion_decision.helpers import (
make_canvas,
make_rows,
make_spec,
)
def _spec(layout):
return make_spec(
layout=layout,
first_position_alignment=FirstPositionAlignment.REQUIRES_PREDECESSOR,
max_canvas=8 if layout is DiffusionLayout.ENCODER_CANVAS else None,
max_context=None,
)
def _canvas(prompt, canvas, positions, name):
return make_canvas(
prompt,
canvas,
positions,
allowed=(1, 2),
question_ids=tuple(f"{name}{i}" for i in range(len(positions))),
)
def _rows():
canvases = (
_canvas((4, 5), (6, 7, 8, 9), (1, 3), "a"),
_canvas((10, 11, 12), (13, 14, 15, 16), (2,), "b"),
)
return make_rows(canvases, weights=(1, 2))
def test_full_sequence_concat_preserves_document_and_label_coordinates():
batch = DecisionTrainingCollator(_spec(DiffusionLayout.FULL_SEQUENCE))(_rows())
assert batch["input_ids"].tolist() == [
[4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
]
assert batch["document_ids"].tolist() == [[0] * 6 + [1] * 7]
assert batch["position_ids"].tolist() == [[0, 1, 2, 3, 4, 5, 0, 1, 2, 3, 4, 5, 6]]
assert batch["decision_label_rows"].tolist() == [[0, 0], [0, -1]]
assert batch["decision_label_positions"].tolist() == [[3, 5], [11, -1]]
assert batch["decision_question_mask"].tolist() == [[True, True], [True, False]]
assert torch.equal(
batch["canvas_update_mask"],
batch["canvas_corruptible_mask"] & ~batch["canvas_input_pinned_mask"],
)
assert [q.position for e in batch["decision_examples"] for q in e.questions] == [
0,
1,
0,
]
assert batch["decision_sources"] == ("source-0", "source-1")
assert batch["decision_slot_counts"].tolist() == [0, 0]
assert batch["decision_draws"] == (None, None)
def test_full_sequence_nested_packing_preserves_documents_labels_and_attention():
rows = _rows()
batch = DecisionTrainingCollator(_spec(DiffusionLayout.FULL_SEQUENCE))(
[[rows[0]], [rows[1]]]
)
assert batch["document_ids"].tolist() == [[0] * 6 + [1] * 7]
assert batch["position_ids"].tolist() == [[0, 1, 2, 3, 4, 5, 0, 1, 2, 3, 4, 5, 6]]
assert batch["decision_label_rows"].tolist() == [[0, 0], [0, -1]]
assert batch["decision_label_positions"].tolist() == [[3, 5], [11, -1]]
from axolotl.core.trainers.diffusion_lm.backends.full_sequence import (
FullSequenceBackend,
)
packed = FullSequenceBackend(mask_token_id=99).pack(
batch["input_ids"],
batch["document_ids"],
batch["semantic_validity"],
batch["position_ids"],
)
attention = packed["attention_mask"]
assert attention[0, 0, 0, 5]
assert attention[0, 0, 6, 12]
assert not attention[0, 0, 0, 6]
assert not attention[0, 0, 6, 5]
def test_multipack_dataloader_covers_each_decision_once():
from axolotl.utils.samplers import MultipackBatchSampler, get_dataset_lengths
rows = [
{**row, "length": len(row["canvas"].prompt_ids) + len(row["canvas"].canvas_ids)}
for row in _rows()
]
dataset = DecisionDataset(rows, {"per_batch_stratified": False})
sampler = MultipackBatchSampler(
SequentialSampler(dataset),
lengths=get_dataset_lengths(dataset),
batch_max_len=20,
batch_size=1,
bin_size=20,
group_size=2,
num_processes=1,
sequential=True,
drop_last=True,
)
batches = list(
DataLoader(
dataset,
batch_sampler=sampler,
collate_fn=DecisionTrainingCollator(_spec(DiffusionLayout.FULL_SEQUENCE)),
)
)
assert len(batches) == 1
assert batches[0]["decision_sources"] == ("source-0", "source-1")
assert batches[0]["document_ids"].tolist() == [[0] * 6 + [1] * 7]
def test_encoder_canvas_keeps_logical_positions_for_postpack_remap():
batch = DecisionTrainingCollator(_spec(DiffusionLayout.ENCODER_CANVAS))(_rows())
logical = batch["diffusion_batch"]
assert logical.encoder_input_ids.tolist() == [[4, 5, 0], [10, 11, 12]]
assert batch["decision_logical_rows"].tolist() == [[0, 0], [1, -1]]
assert batch["decision_label_positions"].tolist() == [[1, 3], [2, -1]]
assert batch["decision_label_rows"].tolist() == [[-1, -1], [-1, -1]]
@pytest.mark.parametrize(
"layout", [DiffusionLayout.FULL_SEQUENCE, DiffusionLayout.ENCODER_CANVAS]
)
def test_collator_preserves_sampled_slot_counts_and_draw_identity(layout):
rows = _rows()
draws = (DecisionDraw(0, 2, 17), DecisionDraw(1, 2, 18))
for row, count, draw in zip(rows, (1, 3), draws, strict=True):
row["decision_slot_count"] = count
row["decision_draw"] = draw
batch = DecisionTrainingCollator(_spec(layout))(rows)
assert batch["decision_slot_counts"].tolist() == [1, 3]
assert batch["decision_draws"] == draws
@pytest.mark.parametrize(
"layout", [DiffusionLayout.FULL_SEQUENCE, DiffusionLayout.ENCODER_CANVAS]
)
def test_collator_emits_prompt_slot_metadata_for_sampled_counts(layout):
rows = _rows()
for row, count, ids in zip(rows, (1, 3), ((7,), (8, 9, 10)), strict=True):
canvas = row["canvas"]
row["canvas"] = DecisionCanvas(
prompt_ids=(*ids, *canvas.prompt_ids),
canvas_ids=canvas.canvas_ids,
label_positions=canvas.label_positions,
allowed_ids=canvas.allowed_ids,
question_ids=canvas.question_ids,
targets=canvas.targets,
pinned_mask=canvas.pinned_mask,
semantic_mask=canvas.semantic_mask,
slot_mask=canvas.slot_mask,
template_length=canvas.template_length,
prompt_slot_mask=(True,) * count + (False,) * len(canvas.prompt_ids),
)
row["decision_slot_count"] = count
row["decision_draw"] = DecisionDraw(count, 0, count)
batch = DecisionTrainingCollator(_spec(layout))(rows)
if layout is DiffusionLayout.FULL_SEQUENCE:
assert batch["decision_prompt_slot_mask"].tolist() == [
[True, False, False, False, False, False, False, True, True, True]
+ [False] * 7
]
assert torch.all(
batch["decision_prompt_slot_mask"] <= batch["canvas_input_pinned_mask"]
)
else:
assert batch["decision_prompt_slot_mask"].tolist() == [
[True, False, False, False, False, False],
[True, True, True, False, False, False],
]
assert torch.equal(
batch["decision_prompt_input_pinned_mask"],
batch["diffusion_batch"].encoder_validity,
)
assert not torch.any(
batch["decision_prompt_slot_mask"]
& batch["diffusion_batch"].encoder_ar_valid_mask
)
assert batch["diffusion_batch"].encoder_ar_valid_mask[0, 1]
def test_builder_compatible_constructor_accepts_standard_collator_options():
collator = DecisionTrainingCollator(
SimpleNamespace(pad_token_id=7),
spec=_spec(DiffusionLayout.FULL_SEQUENCE),
padding=True,
max_length=None,
pad_to_multiple_of=64,
label_pad_token_id=-100,
return_tensors="pt",
)
assert collator.pad_token_id == 7
assert collator.pad_to_multiple_of == 64
def test_full_sequence_rejects_microbatch_above_payload_capacity():
collator = DecisionTrainingCollator(
_spec(DiffusionLayout.FULL_SEQUENCE), physical_payload_capacity=12
)
with pytest.raises(ValueError, match="physical payload capacity"):
collator(_rows())
tests/integrations/diffusion_decision/test_typed_decisions.py
+202
View File
@@ -0,0 +1,202 @@
"""Native multi-question schemas preserve option semantics and soft targets."""
import json
import pytest
import torch
from axolotl.integrations.diffusion_decision.adapters import normalize_record
from axolotl.integrations.diffusion_decision.adapters.typed_decisions import normalize
from axolotl.integrations.diffusion_decision.loss import (
DecisionLabelExample,
DecisionLabelQuestion,
decision_label_loss,
label_target_from_mapping,
)
def sample():
return {
"id": "a",
"state": '{"trace":"observed"}',
"workflow": "trace_review",
"questions": json.dumps(
{
"s": {"type": "score", "criteria": ["low", "high"]},
"c": {
"type": "choice",
"criteria": {"act": "Proceed", "wait": "Delay"},
},
"b": {
"type": "noul",
"criteria": {"true": "Supported", "false": "Refuted"},
},
}
),
"gold": json.dumps(
{
"s": {"probabilities": {"1": 0.8, "0": 0.2}},
"c": {"probabilities": {"wait": 0.4, "act": 0.6}},
"b": {"probabilities": {"false": 1, "true": 0}},
}
),
}
def test_native_multi_question_preserves_all_semantics():
result = normalize(sample())
assert result["questions"]["s"]["levels"] == ["low", "high"]
assert result["labels"]["s"] == {
"kind": "dist",
"candidate_indices": [0, 1],
"candidate_ids": ["0", "1"],
"probs": [0.2, 0.8],
"other_probability": 0.0,
}
assert result["questions"]["c"]["options"] == [
{"name": "act", "description": "Proceed"},
{"name": "wait", "description": "Delay"},
]
assert result["labels"]["c"]["probs"] == [0.6, 0.4]
assert result["labels"]["b"] == {
"kind": "dist",
"candidate_indices": [1],
"candidate_ids": ["no"],
"probs": [1.0],
"other_probability": 0.0,
}
assert result["questions"]["b"]["criteria"]["true"] == "Supported"
assert result["family"] == "trace_review"
assert result["state"] == {"trace": "observed"}
assert result["source_metadata"]["provenance"] == {
"dataset": "LocalLLaMA/typed-decisions",
"workflow": "trace_review",
"split": None,
"row_split": None,
"source_id": "a",
}
def test_missing_probability_rejected_without_inventing_mass():
row = sample()
gold = json.loads(row["gold"])
del gold["s"]["probabilities"]["0"]
row["gold"] = json.dumps(gold)
with pytest.raises(ValueError, match="each alternative"):
normalize(row)
def test_more_than_four_soft_alternatives_preserve_full_distribution():
row = sample()
questions = json.loads(row["questions"])
gold = json.loads(row["gold"])
questions["s"]["criteria"] = ["a", "b", "c", "d", "e"]
gold["s"]["probabilities"] = {"0": 0.1, "1": 0.2, "2": 0.3, "3": 0.15, "4": 0.25}
row["questions"] = json.dumps(questions)
row["gold"] = json.dumps(gold)
target = normalize(row)["labels"]["s"]
assert target == {
"kind": "dist",
"candidate_indices": [0, 1, 2, 3, 4],
"candidate_ids": ["0", "1", "2", "3", "4"],
"probs": [0.1, 0.2, 0.3, 0.15, 0.25],
"other_probability": 0.0,
}
result = decision_label_loss(
torch.zeros((1, 1, 5)),
[
DecisionLabelExample(
(
DecisionLabelQuestion(
0,
(0, 1, 2, 3, 4),
label_target_from_mapping(target),
),
)
)
],
torch.tensor([[True]]),
label_softmax="both",
)
assert torch.isfinite(result.loss)
def test_missing_question_target_rejected():
row = sample()
row["gold"] = "{}"
with pytest.raises(ValueError, match="exactly one target"):
normalize(row)
def test_training_is_fail_closed_to_declared_train_split():
result = normalize_record(
"diffusion_decision.typed_decisions",
sample(),
training=True,
source_split="train",
)
assert result["source_metadata"]["provenance"]["split"] == "train"
with pytest.raises(ValueError, match="declared train split"):
normalize_record(
"diffusion_decision.typed_decisions",
sample(),
training=True,
source_split="test",
)
def test_row_split_mismatch_is_rejected():
row = sample()
row["split"] = "test"
with pytest.raises(ValueError, match="disagrees"):
normalize_record(
"diffusion_decision.typed_decisions",
row,
training=True,
source_split="train",
)
def test_native_hard_smoothing_preserves_question_semantics():
row = sample()
gold = json.loads(row["gold"])
gold["b"]["smoothing"] = 0.1
row["gold"] = json.dumps(gold)
baseline = normalize(sample())
result = normalize(row)
assert result["questions"] == baseline["questions"]
assert result["labels"]["s"] == baseline["labels"]["s"]
assert result["labels"]["c"] == baseline["labels"]["c"]
assert result["labels"]["b"] == {"kind": "hard", "gold_idx": 1, "smoothing": 0.1}
@pytest.mark.parametrize("smoothing", [True, -0.1, 1.0, float("nan")])
def test_native_hard_smoothing_validates_metadata(smoothing):
row = sample()
gold = json.loads(row["gold"])
gold["b"]["smoothing"] = smoothing
row["gold"] = json.dumps(gold)
with pytest.raises(ValueError, match="hard smoothing"):
normalize(row)
def test_native_hard_smoothing_rejects_soft_targets():
row = sample()
gold = json.loads(row["gold"])
gold["s"]["smoothing"] = 0.1
row["gold"] = json.dumps(gold)
with pytest.raises(ValueError, match="one-hot"):
normalize(row)
def test_native_hard_smoothing_uses_one_hot_probability_not_gold_label():
row = sample()
gold = json.loads(row["gold"])
gold["b"]["label"] = "true"
gold["b"]["smoothing"] = 0.1
row["gold"] = json.dumps(gold)
assert normalize(row)["labels"]["b"] == {
"kind": "hard",
"gold_idx": 1,
"smoothing": 0.1,
}
tests/integrations/diffusion_decision/test_vocab_size_lookup.py
+64
View File
@@ -0,0 +1,64 @@
"""Avoid expensive tokenizer vocabulary scans when model metadata is available."""
from types import SimpleNamespace
from axolotl.integrations.diffusion_decision import datasets
def test_canvas_row_uses_model_vocabulary_without_counting_tokenizer(monkeypatch):
class Tokenizer:
eos_token_id = 11
pad_token_id = 11
def __len__(self):
raise AssertionError("unexpected full tokenizer vocabulary count")
captured = {}
def build(_tokenizer, _record, **kwargs):
captured.update(kwargs)
return SimpleNamespace(prompt_ids=(1,), canvas_ids=(100,) * 128)
monkeypatch.setattr(datasets, "build_decision_canvas", build)
monkeypatch.setattr(datasets, "permute_record", lambda record, **kwargs: record)
monkeypatch.setattr(
datasets, "decision_example_from_canvas", lambda *args, **kwargs: None
)
cfg = {
"model_config": {"vocab_size": 131072},
"diffusion_lm": {"mask_token_id": 100},
"diffusion_decision": {},
}
row = datasets._canvas_row(
Tokenizer(),
{"source": "test"},
cfg,
1.0,
spec=SimpleNamespace(max_canvas=128, noise=datasets.DiffusionNoise.ABSORBING),
)
assert captured["vocab_size"] == 131072
assert row["length"] == 129
def test_missing_model_metadata_counts_vocabulary_once_per_pool(monkeypatch):
class Tokenizer:
calls = 0
def __len__(self):
self.calls += 1
return 131072
tokenizer = Tokenizer()
seen = []
def build_row(_tokenizer, record, cfg, weight, **kwargs):
seen.append(kwargs["vocab_size"])
return record
monkeypatch.setattr(datasets, "_canvas_row", build_row)
rows = [{"source": "test", "id": str(index)} for index in range(20)]
actual, dropped = datasets._canvas_rows(tokenizer, rows, {}, {}, None)
assert actual == rows
assert dropped == 0
assert seen == [131072] * 20
assert tokenizer.calls == 1
tests/test_native_diffusion_model_support.py
+50
View File
@@ -819,6 +819,56 @@ def test_trainable_token_adapter_routes_to_efficient_merge(monkeypatch, tmp_path
adapter_has_trainable_token_deltas(missing_delta)
def test_merged_export_preserves_valid_decision_manifest(tmp_path, monkeypatch):
from collections import OrderedDict
from types import SimpleNamespace
from axolotl.integrations.base import PluginManager
from axolotl.integrations.diffusion_decision.manifest import build_decision_manifest
from axolotl.integrations.diffusion_decision.plugin import DiffusionDecisionPlugin
adapter = tmp_path / "adapter"
adapter.mkdir()
cfg = DictDefault(
{
"base_model": "nvidia/Nemotron-Labs-Diffusion-3B",
"revision_of_model": "0d51902da1f8869f83413ce642fab402fa5641e0",
"model_config_type": "nemotron_labs_diffusion",
"output_dir": str(tmp_path),
"diffusion_lm": {"canvas_width": 128, "mask_token_id": 100},
"diffusion_decision": {
"layout": "prompt_slots",
"latent": {"mode": "none", "num_slots": 0},
},
}
)
tokenizer = SimpleNamespace(
bos_token_id=1,
eos_token_id=2,
pad_token_id=0,
unk_token_id=3,
mask_token_id=100,
all_special_ids=[0, 1, 2, 3, 100],
)
model = SimpleNamespace(
config=SimpleNamespace(
mask_token_id=100,
_commit_hash="0d51902da1f8869f83413ce642fab402fa5641e0",
)
)
manifest = build_decision_manifest(cfg, tokenizer=tokenizer, model=model)
source = manifest.write_to(adapter)
merged = tmp_path / "merged"
merged.mkdir()
monkeypatch.setattr(
PluginManager.get_instance(),
"plugins",
OrderedDict({"diffusion_decision": DiffusionDecisionPlugin()}),
)
PluginManager.get_instance().post_lora_merge(cfg, str(adapter), str(merged))
assert (merged / source.name).read_bytes() == source.read_bytes()
def test_native_generation_dispatches_and_preserves_legacy_seam():
from types import SimpleNamespace