diff --git a/.agents/skills/areno-build-singleturn-game-demo/SKILL.md b/.agents/skills/areno-build-singleturn-game-demo/SKILL.md
new file mode 100644
index 00000000..d7670299
--- /dev/null
+++ b/.agents/skills/areno-build-singleturn-game-demo/SKILL.md
@@ -0,0 +1,98 @@
+---
+name: areno-build-singleturn-game-demo
+description: Autonomously research, implement, train, evaluate, visualize, and submit a genuinely new AReno RLVR mini-game whose mechanically generated samples are independent state-to-single-output tasks. Use when the user requests a complete new game demo with real reward improvement, not for multi-turn agent workflows, SFT-only examples, or design-only proposals.
+---
+
+# Build an AReno Single-Turn Game Demo
+
+Deliver a working, reproducible AReno mini-game rather than a proposal. Continue
+through research, implementation, data generation, baseline evaluation, RLVR
+training, held-out evaluation, WebUI delivery, and PR submission until the
+evidence meets the requested acceptance criteria or a genuine external blocker
+remains.
+
+## Non-negotiable task shape
+
+Mechanically generable sequential games must be transformed into independent
+single-turn examples:
+
+```text
+complete state -> exactly one generation -> deterministic reward
+```
+
+- Do not use a multi-turn conversation, agent loop, tool-feedback loop, or a
+  sequence of assistant/tool messages for training or evaluation.
+- Generate intermediate states directly when the source game is sequential.
+- Each prompt contains the complete public state and output contract, but no
+  oracle answer, private reward field, or reversible instance identifier.
+- Each reward call scores only that prompt and its one completion.
+- A WebUI may show successive states for human play; this does not change the
+  independent single-turn training contract.
+
+## Before acting
+
+1. Read repository `AGENTS.md`, `CODEMAP.md`, current CLI help, and the
+   repository-local training, serving, correctness, and capacity skills that
+   match the requested work.
+2. Inspect current `examples/agentic/`, `examples/sft/`, `examples/vl/`, and
+   any other example directories completely enough to inventory every demo's
+   game/task, state representation, output, reward target, and training paradigm.
+3. Read [research-and-selection.md](references/research-and-selection.md) before
+   choosing a game.
+4. Read [implementation-contract.md](references/implementation-contract.md)
+   before writing source or data.
+5. Read [experiment-runbook.md](references/experiment-runbook.md) before running
+   dataset inspection, baseline, training, evaluation, or checkpoint commands.
+6. Read [webui-and-reporting.md](references/webui-and-reporting.md) before
+   implementing the UI, writing README results, or opening the PR.
+
+Use latest `origin/main` in a dedicated branch/worktree and preserve unrelated
+user changes. Verify every interface against checked-out source rather than
+memory. Do not modify public config or CLI surfaces unless the request explicitly
+requires it and repository policy permits it.
+
+## Authorization boundary
+
+This skill does not itself authorize paid compute, remote mutations, public
+services, pushes, or PRs. Treat an invocation that explicitly requests autonomous
+training, serving, pushing, and PR creation as authorization for those in-scope
+actions. Otherwise obtain only the missing authorization when it becomes
+necessary. Never ask the user to perform routine steps the agent can safely do.
+
+## Required execution loop
+
+1. Inventory existing demos and record overlap risks.
+2. Research at least five viable games using accessible Chinese sources and
+   cross-check selected rules with two independent sources.
+3. Select a game only if it has a deterministic oracle, scalable mechanical
+   generation, strong random-versus-reasoning separation, and a clear future UI.
+4. Implement the decoupled game/oracle, public state view, generator, loader,
+   prompt contract, parser, reward, evaluation, and focused tests.
+5. Generate leak-free train/validation/test data and run legality/oracle self-checks.
+6. Inspect normalized data before training. Run a bounded smoke workload.
+7. Evaluate the unmodified base checkpoint on fixed held-out splits.
+8. Run real single-turn RLVR training. Diagnose and iterate when learning is weak;
+   do not change the held-out set or hide failed runs.
+9. Evaluate checkpoints under matched conditions and stop only after the target
+   improvement is achieved and reward reaches a measured plateau.
+10. Implement and probe a polished playable WebUI, then start requested model and
+    UI services if authorized.
+11. Complete documentation, remove generated artifacts from version control,
+    review the diff, commit, push, and open a focused PR.
+
+## Completion gate
+
+Do not claim success unless actual logs and held-out evaluation establish all of
+the following:
+
+- mean reward improves by at least `0.15` absolute or `30%` relative;
+- core success/accuracy clearly improves, not merely formatting validity;
+- at least two evaluation seeds agree in direction;
+- train/validation/test remain isolated and baseline/post-training conditions match;
+- focused tests, smoke validation, real training, checkpoint save/reload as
+  applicable, serving probe, and WebUI probe have actually run;
+- README contains commands, all formal evaluation runs, failed experiments,
+  known limitations, source links, demo-difference evidence, and WebUI design.
+
+If available compute cannot meet the gate, report the exact blocker and evidence;
+never fabricate a curve, checkpoint, test result, or service status.
diff --git a/.agents/skills/areno-build-singleturn-game-demo/agents/openai.yaml b/.agents/skills/areno-build-singleturn-game-demo/agents/openai.yaml
new file mode 100644
index 00000000..ee246d6b
--- /dev/null
+++ b/.agents/skills/areno-build-singleturn-game-demo/agents/openai.yaml
@@ -0,0 +1,4 @@
+interface:
+  display_name: "Build an AReno Single-Turn Game Demo"
+  short_description: "Research, train, evaluate, and ship RLVR games"
+  default_prompt: "Use $areno-build-singleturn-game-demo to autonomously build and validate a new single-turn RLVR game demo."
diff --git a/.agents/skills/areno-build-singleturn-game-demo/references/experiment-runbook.md b/.agents/skills/areno-build-singleturn-game-demo/references/experiment-runbook.md
new file mode 100644
index 00000000..50b4a80e
--- /dev/null
+++ b/.agents/skills/areno-build-singleturn-game-demo/references/experiment-runbook.md
@@ -0,0 +1,104 @@
+# RLVR experiment runbook
+
+Read this reference before data inspection, model inference, training,
+evaluation, or checkpoint handling.
+
+## Workflow defaults
+
+- Do not hardcode a model family or checkpoint in the skill. Use the checkpoint
+  explicitly requested by the invoking user. If none is specified, select a
+  currently supported AReno checkpoint that fits the available compute and task,
+  then record the exact choice and rationale before baseline evaluation.
+- Training paradigm: RLVR only. Never replace it with SFT and never fabricate an
+  agentic loop. Select the current supported single-turn rollout RL algorithm
+  from checked-out AReno API and document the choice.
+- Optimizer: always use AReno's current Adam4bit option (expected CLI spelling
+  `--adam-4bit`; verify it from current help before use).
+- Never reinstall AReno. Use the installed package, environment, and repository
+  mechanisms already available.
+- Never use eager decode. Preserve the normal CUDA-graph decode path and fix or
+  tune the actual issue instead of disabling graphs.
+- Save experiment checkpoints and logs below a task-specific directory under
+  `/new/`, which is a mounted path. Never place them in the repository.
+- Retain only the latest useful checkpoint during training. Prefer the current
+  supported keep-latest option; otherwise delete only verified old checkpoints
+  inside the task-specific checkpoint directory immediately before a new save.
+
+## Environment and data gate
+
+Before building the training command:
+
+1. Record commit, branch, `areno env --json`, `areno check`, GPU topology and
+   memory, installed AReno and model paths, and model hub selection.
+2. Read current `areno train --help` and repository training or capacity skills.
+3. Inspect raw and normalized train, validation, and test samples using repository
+   dataset inspection tools. Do not train until inspection reports success and
+   confirms one prompt or messages input per row with no oracle leakage.
+4. Run generator self-checks and compare canonical hashes across splits.
+
+Use ModelScope for remote AReno assets when repository policy requires it. Do
+not silently switch hubs.
+
+## Baseline protocol
+
+Evaluate the untouched base checkpoint on fixed held-out validation and test
+data. Record:
+
+- mean reward;
+- task success or accuracy;
+- legal or parseable output rate;
+- results by difficulty bucket;
+- dataset size and split seed;
+- model or checkpoint, commit, temperature, max tokens, and all inference settings;
+- random-policy lower bound and oracle upper bound when applicable.
+
+Use at least two formal evaluation seeds. Save machine-readable output outside
+the repository and a compact checked-in result summary without private data.
+
+## Training and capacity loop
+
+Start with a small real smoke workload that exercises rollout, reward, backward,
+optimizer step, and checkpoint save. Then run enough real RLVR steps to observe
+a learning curve.
+
+Capacity tuning order:
+
+1. If rollout OOMs, reduce `max_running_prompts` before changing semantic context
+   or generation length.
+2. If training OOMs, reduce `mini_bs` before semantic token limits.
+3. Respect `batch_size * n_samples` total demand and keep concurrency separate.
+4. Do not use eager decode as an OOM workaround.
+
+When reward grows slowly, diagnose prompt clarity, parsing, reward density,
+difficulty mix, sampling, batch or group size, and training budget. Consider
+learning rate `1e-5` with minimum learning rate `1e-6` after evidence indicates
+the original schedule is too weak. Do not tune on held-out test data.
+
+Evaluate intermediate checkpoints under the exact baseline protocol. Stop
+training only after:
+
+- the acceptance improvement is met; and
+- held-out mean reward improves by less than `0.02` across two consecutive
+  scheduled evaluations, or another predeclared statistically comparable
+  plateau rule.
+
+Do not stop merely because training reward is high. Do not continue indefinitely
+after held-out reward has clearly plateaued.
+
+## Post-training comparison
+
+Use the same held-out rows, parser, reward, temperature, decoding settings, and
+difficulty buckets as baseline. Success requires:
+
+- at least `0.15` absolute mean-reward gain or `30%` relative gain;
+- clear core correctness gain;
+- consistent direction on at least two eval seeds;
+- evidence the gain is not only formatting compliance.
+
+Report every formal run, including failed configurations. Check for answer
+leakage, position bias, generator shortcuts, memorized seed mappings, duplicated
+states, and format-only learning.
+
+If the target remains unmet after reasonable prompt, reward, curriculum, and
+optimizer iteration, replace the game with a more trainable researched candidate
+rather than misrepresenting the result.
diff --git a/.agents/skills/areno-build-singleturn-game-demo/references/implementation-contract.md b/.agents/skills/areno-build-singleturn-game-demo/references/implementation-contract.md
new file mode 100644
index 00000000..cfb2841e
--- /dev/null
+++ b/.agents/skills/areno-build-singleturn-game-demo/references/implementation-contract.md
@@ -0,0 +1,101 @@
+# Implementation contract
+
+Read this reference before implementing the game, dataset, reward, or tests.
+
+## Repository-first design
+
+Inspect current dataset loader and generator contracts, reward API, trainer, CLI,
+evaluation path, and the two or three closest high-quality demos. Follow current
+public API and local style. Do not copy an obsolete example or introduce an
+agentic API merely because the source game is sequential.
+
+Create a self-contained example in the most appropriate current `examples/`
+location. File names may follow current conventions, but responsibilities must be
+obvious. The demo normally needs equivalents of:
+
+```text
+README.md
+dataset_generator.py
+dataset_loader.py
+game.py and/or solver.py
+reward.py
+single-turn inference entry
+eval.py when common evaluation is insufficient
+web_ui.py after training evidence is complete
+```
+
+## Game and public-view boundary
+
+- Represent internal state as stable JSON-compatible data.
+- Define a finite, structured, independently verifiable action or full solution.
+- Decouple game rules, transition or oracle logic, reward, prompt rendering, and UI.
+- Provide a pure function or equivalent method that converts internal state to a
+  public view.
+- Never expose oracle answers, reward-private labels, seeds, or hidden metadata in
+  the public view.
+- Do not make ANSI codes, terminal coordinates, rendered prompt text, or CLI output
+  the canonical state.
+
+## Dataset generator
+
+The generator must:
+
+- expose CLI parameters for seed, sample counts, difficulty range, and output path;
+- use a local reproducible RNG rather than global randomness;
+- create distinct train, validation, and test splits;
+- use disjoint seeds or non-overlapping IDs and prevent duplicate or equivalent
+  instances across splits;
+- generate complete independent states, never conversation trajectories;
+- store oracle answer, difficulty, and needed private metadata outside the visible
+  prompt;
+- support both fast smoke data and formal data;
+- validate legality, solver correctness, output reproducibility, split isolation,
+  and any canonical-equivalence rule before reporting success.
+
+Do not commit formal generated datasets or checkpoints unless the repository
+explicitly tracks small fixtures. Tiny deterministic test fixtures are acceptable.
+
+## Prompt and output
+
+Use a compact, self-contained prompt:
+
+```text
+Input: complete public game state
+Output: exactly one constrained action or complete structured solution
+```
+
+- State the minimum rules needed to solve the instance.
+- Require one generation and a minimal stable format.
+- Do not request chain-of-thought.
+- Do not include oracle values, hidden fields, or reversible ID-to-answer mappings.
+- Do not use assistant or tool history, environment feedback, or future observations.
+- Keep training and evaluation rendering identical unless a documented experiment
+  deliberately compares them.
+
+## Deterministic anti-exploit reward
+
+The reward must parse safely and return a score rather than raising on malformed
+output. Define signals for:
+
+- fully correct or optimal output: maximum score;
+- legal but suboptimal output where the game permits it;
+- partially correct progress when a meaningful verifier-derived dense signal exists;
+- illegal action;
+- malformed output.
+
+Highest reward must mean the task is genuinely solved. Defend against multiple
+answers, prompt copying, extra prose, overly long output, duplicate actions,
+injection-like suffixes, NaN or overflow values, and any alternative syntax that
+could bypass validation. Do not call another LLM and do not use private labels to
+grant unearned reward.
+
+Add focused CPU-safe tests for:
+
+- game transitions and edge states;
+- solver or oracle correctness and determinism;
+- generator repeatability and split leakage detection;
+- loader normalization and prompt privacy;
+- correct, wrong, partial, illegal, malformed, oversized, multi-answer, and
+  adversarial completions;
+- public-view privacy and JSON serialization;
+- evaluation aggregation and difficulty buckets.
diff --git a/.agents/skills/areno-build-singleturn-game-demo/references/research-and-selection.md b/.agents/skills/areno-build-singleturn-game-demo/references/research-and-selection.md
new file mode 100644
index 00000000..20cf7fcf
--- /dev/null
+++ b/.agents/skills/areno-build-singleturn-game-demo/references/research-and-selection.md
@@ -0,0 +1,76 @@
+# Research and game selection
+
+Read this reference before choosing the game or creating its directory.
+
+## Existing-demo exclusion audit
+
+Inspect the current branch, not a remembered list. At minimum inventory every
+demo under `examples/agentic/`, `examples/sft/`, `examples/vl/`, multimodal,
+math, and any newly added example roots. For each record:
+
+- task/game name;
+- core state representation;
+- model output form;
+- reward objective;
+- training paradigm.
+
+Reject a candidate whose core mechanic or task structure overlaps an existing
+demo. A new story, board size, symbols, wording, or visual theme is not a new
+mechanic. In particular, do not recreate:
+
+- Tic-Tac-Toe, Gomoku, or other line-making move games;
+- Codebreaker, Mastermind, terminal likeness deduction, or isomorphic code games;
+- DuelGrid-style board combat or local tactical placement;
+- shopping or product constraint selection;
+- coding, math verification, or music-generation examples with a cosmetic theme;
+- anything newly present in checked-out `examples/` even if absent above.
+
+The final README must include a concrete “Differences from existing demos” table
+based on mechanics, state, output, reward, and training shape.
+
+## Chinese-first internet research
+
+Compare at least five candidate mini-games before selecting one. Prefer public or
+well-known classic logic games whose rules can be independently reimplemented.
+
+The environment may lack overseas network access. Search in Chinese and prefer
+mainland-accessible sources such as Baidu or 360 search and encyclopedias, Zhihu,
+Bilibili, CSDN, CNBlogs, Juejin, Jianshu, Chinese university or education sites,
+and Chinese game-rule sites. GitHub may be used for necessary API or open-source
+implementation references, but not as the only rules source.
+
+- Do not depend on an overseas search engine, community, video site, or rule site.
+- If a page fails, times out, requires a proxy, or cannot be fetched reliably,
+  switch sources immediately instead of repeatedly waiting.
+- Search snippets may support initial screening, but before final selection open
+  at least one full Chinese rules page.
+- Cross-check selected rules using at least two independent accessible Chinese
+  sources.
+- Keep actual working URLs in README and summarize rules in original words; do
+  not copy long passages.
+
+## Candidate requirements
+
+A candidate must satisfy every item:
+
+1. State generation is deterministic, cheap, and scalable.
+2. A reliable algorithm computes valid or optimal answers.
+3. Reward is deterministic and needs no human or LLM judge.
+4. One input contains everything needed for one output.
+5. Difficulty is parameterized and supports diverse data.
+6. Random guessing is materially worse than genuine reasoning.
+7. A realistic training budget can expose a learning curve.
+8. It does not overlap current AReno demos.
+9. It uses no copied proprietary puzzle bank, prose, or art.
+10. JSON state and finite structured actions naturally support an intuitive WebUI.
+
+## Required comparison record
+
+Put a table in README with at least these columns:
+
+| Candidate | Rules/mechanic | Single-turn conversion | Oracle | Reward | Scale/difficulty | Existing-demo similarity risk | UI fit | Decision |
+| --- | --- | --- | --- | --- | --- | --- | --- | --- |
+
+Explain the final selection using trainability and validation quality, not novelty
+alone. Preserve sources actually used, including sources for rejected candidates
+when they materially informed the decision.
diff --git a/.agents/skills/areno-build-singleturn-game-demo/references/webui-and-reporting.md b/.agents/skills/areno-build-singleturn-game-demo/references/webui-and-reporting.md
new file mode 100644
index 00000000..d2d30f01
--- /dev/null
+++ b/.agents/skills/areno-build-singleturn-game-demo/references/webui-and-reporting.md
@@ -0,0 +1,88 @@
+# WebUI, documentation, and delivery
+
+Read this reference before implementing WebUI, finalizing README, starting
+services, or opening a PR.
+
+## Playable WebUI
+
+After training and evaluation evidence is complete, implement a polished playable
+WebUI using current Tic-Tac-Toe or terminal-hacking WebUI code as engineering
+references, not as visual or gameplay templates.
+
+- Render from public JSON view, never by parsing the training prompt.
+- Use clear board, card, node, path, counter, or other components appropriate to
+  the selected mechanic, with responsive layout, deliberate visual hierarchy,
+  loading, error, win or loss, and disabled states.
+- Map every UI interaction to the same structured action schema verified by the
+  game and reward layer.
+- Keep display or session progression outside training data. Human play may be
+  sequential while model inference remains one independent state-to-output request.
+- Handle model latency asynchronously without hiding the current game state.
+- Add a representative public-view JSON example to README.
+
+If authorized by the invoking request, start the trained checkpoint through
+AReno's OpenAI-compatible server on port `8000` and WebUI on port `8001`. Use
+the current serve command, keep CUDA graphs enabled, probe `/v1/models` and a
+real completion, then probe UI and gameplay API. Report actual process IDs, log
+paths, and whether services remain running.
+
+## README contents
+
+The demo README must contain:
+
+1. rules and accessible Chinese source links;
+2. comparison of at least five researched candidates;
+3. differences from every current overlapping-risk demo;
+4. single-turn state-to-output contract;
+5. internal state, public view, action schema, oracle, and reward design;
+6. generator CLI, split sizes, seeds, difficulty, leak checks, and data inspection;
+7. smoke, baseline, training, checkpoint, post-eval, serving, and WebUI commands;
+8. full experiment configuration and baseline or post-training table;
+9. per-seed and per-difficulty results;
+10. all failed formal experiments and resulting iterations;
+11. tests actually run;
+12. Future WebUI visualization layout, components, rendering, and action mapping;
+13. known limitations and next steps.
+
+Do not claim any command, result, checkpoint, or service that was not observed
+in the current work.
+
+## Final repository hygiene
+
+Before committing:
+
+- inspect `git diff` and status;
+- remove caches, local environments, downloaded models, generated formal datasets,
+  checkpoints, logs, credentials, endpoints, and absolute local paths;
+- preserve only source, focused tests, tiny fixtures if justified, README, and
+  compact result records;
+- verify no unrelated user changes are included;
+- run relevant formatting, linting, and focused tests;
+- verify documentation commands against current help.
+
+Commit on a dedicated branch, push, and open a focused PR only when authorized.
+Link the motivating issue if one exists, state exact validation, and disclose
+skipped GPU or platform checks. Do not auto-close an issue whose acceptance
+evidence remains incomplete.
+
+## Final response contract
+
+Report all of the following with direct artifact or PR links where available:
+
+1. selected game and rationale;
+2. non-overlap evidence;
+3. single-turn state or output design;
+4. files added or changed;
+5. dataset sizes, splits, seeds, and difficulty distribution;
+6. reward formula and exploit resistance;
+7. baseline versus post-training table;
+8. results for at least two eval seeds;
+9. exact training and evaluation commands;
+10. focused test results;
+11. failed experiments and iterations;
+12. WebUI state schema, layout, and interaction mapping;
+13. remaining limitations and recommended next step;
+14. running service endpoints and probe status when requested.
+
+Evidence takes priority over narrative. If blocked, identify the missing
+permission, resource, credential, or external state and the last verified milestone.
diff --git a/areno/accel/__init__.py b/areno/accel/__init__.py
index bcd806c0..590dad00 100644
--- a/areno/accel/__init__.py
+++ b/areno/accel/__init__.py
@@ -30,7 +30,12 @@
 from areno.accel.linear import areno_grouped_linear, areno_linear
 from areno.accel.moe import areno_moe_permute, areno_moe_topk_permute, areno_moe_unpermute
 from areno.accel.normalization import areno_optional_scale_rmsnorm, areno_rmsnorm, areno_rmsnorm_silu_gate
-from areno.accel.optimizer import areno_adamw_4bit_step, areno_adamw_8bit_step, areno_adamw_fp32_master_step
+from areno.accel.optimizer import (
+    areno_adamw_4bit_step,
+    areno_adamw_8bit_step,
+    areno_adamw_fp32_master_step,
+    areno_adamw_fp32_state_step,
+)
 from areno.accel.router import areno_grouped_topk_router
 from areno.accel.routing import areno_moe_align
 from areno.accel.topk import areno_topk_softmax
@@ -53,6 +58,7 @@
     "areno_adamw_4bit_step",
     "areno_adamw_8bit_step",
     "areno_adamw_fp32_master_step",
+    "areno_adamw_fp32_state_step",
     "areno_optional_scale_rmsnorm",
     "areno_rmsnorm",
     "areno_rmsnorm_silu_gate",
diff --git a/areno/accel/csrc/extension.cpp b/areno/accel/csrc/extension.cpp
index 494233b8..20d7f689 100644
--- a/areno/accel/csrc/extension.cpp
+++ b/areno/accel/csrc/extension.cpp
@@ -209,6 +209,8 @@ void areno_adamw_8bit_step_cuda(
     torch::Tensor exp_avg_scale,
     torch::Tensor exp_avg_sq_q,
     torch::Tensor exp_avg_sq_scale,
+    torch::Tensor signed_codebook,
+    torch::Tensor unsigned_codebook,
     int64_t quant_block_size,
     double beta1,
     double beta2,
@@ -217,11 +219,24 @@ void areno_adamw_8bit_step_cuda(
     double eps,
     double step_size,
     double bias_correction2_sqrt);
+void areno_adamw_fp32_state_step_cuda(
+    torch::Tensor model,
+    torch::Tensor grad,
+    torch::Tensor exp_avg,
+    torch::Tensor exp_avg_sq,
+    double beta1,
+    double beta2,
+    double effective_lr,
+    double weight_decay,
+    double eps,
+    double step_size,
+    double bias_correction2_sqrt);
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("areno_adamw_fp32_master_step", &areno_adamw_fp32_master_step_cuda, "ARENO compact FP32-master AdamW step");
   m.def("areno_adamw_4bit_step", &areno_adamw_4bit_step_cuda, "ARENO packed block-wise AdamW4bit step");
   m.def("areno_adamw_8bit_step", &areno_adamw_8bit_step_cuda, "ARENO block-wise 8-bit AdamW step");
+  m.def("areno_adamw_fp32_state_step", &areno_adamw_fp32_state_step_cuda, "ARENO FP32-state AdamW step");
   m.def("areno_silu_and_mul", &areno_silu_and_mul_cuda, "ARENO SiLU and multiply");
   m.def("areno_gelu_tanh_and_mul", &areno_gelu_tanh_and_mul_cuda, "ARENO tanh GELU and multiply");
   m.def("areno_silu", &areno_silu_cuda, "ARENO SiLU");
diff --git a/areno/accel/csrc/optimizer.cu b/areno/accel/csrc/optimizer.cu
index 4c885b56..85e38f31 100644
--- a/areno/accel/csrc/optimizer.cu
+++ b/areno/accel/csrc/optimizer.cu
@@ -31,6 +31,25 @@ __device__ __forceinline__ uint8_t nearest_signed_dynamic_code(float normalized)
   return best;
 }
 
+__device__ __forceinline__ uint8_t nearest_dynamic_code(float value, const float* codebook) {
+  int lower = 0;
+  int upper = 255;
+  while (lower < upper) {
+    const int middle = (lower + upper) >> 1;
+    if (codebook[middle] < value) {
+      lower = middle + 1;
+    } else {
+      upper = middle;
+    }
+  }
+  if (lower == 0) {
+    return 0;
+  }
+  const float left_distance = fabsf(value - codebook[lower - 1]);
+  const float right_distance = fabsf(codebook[lower] - value);
+  return static_cast<uint8_t>(left_distance <= right_distance ? lower - 1 : lower);
+}
+
 __device__ __forceinline__ float adamw_update(
     float master,
     float grad,
@@ -397,6 +416,8 @@ __global__ void adamw_8bit_blockwise_kernel(
     float* exp_avg_scale,
     uint8_t* exp_avg_sq_q,
     float* exp_avg_sq_scale,
+    const float* signed_codebook,
+    const float* unsigned_codebook,
     int64_t numel,
     int64_t quant_block_size,
     float beta1,
@@ -410,6 +431,7 @@ __global__ void adamw_8bit_blockwise_kernel(
   constexpr int max_warps = 8;
   __shared__ float warp_moment_maxima[max_warps];
   __shared__ float warp_variance_maxima[max_warps];
+  __shared__ int invalid_block;
 
   const int64_t block_start = static_cast<int64_t>(blockIdx.x) * quant_block_size;
   const int64_t remaining = numel - block_start;
@@ -418,12 +440,16 @@ __global__ void adamw_8bit_blockwise_kernel(
   const float old_variance_scale = exp_avg_sq_scale[blockIdx.x];
   float local_moment_max = 0.0f;
   float local_variance_max = 0.0f;
+  if (threadIdx.x == 0) {
+    invalid_block = 0;
+  }
+  __syncthreads();
 
   for (int64_t offset = threadIdx.x; offset < block_numel; offset += blockDim.x) {
     const int64_t index = block_start + offset;
     const float gradient = load_grad(grad, index);
-    float moment = (static_cast<int>(exp_avg_q[index]) - 128) * old_moment_scale;
-    float variance = static_cast<float>(exp_avg_sq_q[index]) * old_variance_scale;
+    float moment = signed_codebook[exp_avg_q[index]] * old_moment_scale;
+    float variance = unsigned_codebook[exp_avg_sq_q[index]] * old_variance_scale;
     float weight = load_model(model, index);
     weight = adamw_update(
         weight,
@@ -437,7 +463,9 @@ __global__ void adamw_8bit_blockwise_kernel(
         eps,
         step_size,
         bias_correction2_sqrt);
-    store_model(model, index, weight);
+    if (!isfinite(gradient) || !isfinite(moment) || !isfinite(variance) || !isfinite(weight)) {
+      atomicExch(&invalid_block, 1);
+    }
     local_moment_max = fmaxf(local_moment_max, fabsf(moment));
     local_variance_max = fmaxf(local_variance_max, variance);
   }
@@ -454,6 +482,9 @@ __global__ void adamw_8bit_blockwise_kernel(
     warp_variance_maxima[warp] = local_variance_max;
   }
   __syncthreads();
+  if (invalid_block != 0) {
+    return;
+  }
 
   if (warp == 0) {
     const int warp_count = blockDim.x / warp_size;
@@ -466,8 +497,8 @@ __global__ void adamw_8bit_blockwise_kernel(
           fmaxf(block_variance_max, __shfl_down_sync(0xFFFFFFFFu, block_variance_max, offset));
     }
     if (lane == 0) {
-      exp_avg_scale[blockIdx.x] = fmaxf(block_moment_max / 127.0f, 1.0e-30f);
-      exp_avg_sq_scale[blockIdx.x] = fmaxf(block_variance_max / 255.0f, 1.0e-30f);
+      exp_avg_scale[blockIdx.x] = block_moment_max;
+      exp_avg_sq_scale[blockIdx.x] = block_variance_max;
     }
   }
   __syncthreads();
@@ -476,14 +507,49 @@ __global__ void adamw_8bit_blockwise_kernel(
   for (int64_t offset = threadIdx.x; offset < block_numel; offset += blockDim.x) {
     const int64_t index = block_start + offset;
     const float gradient = load_grad(grad, index);
-    const float moment =
-        beta1 * (static_cast<int>(exp_avg_q[index]) - 128) * old_moment_scale + (1.0f - beta1) * gradient;
-    const float variance = beta2 * static_cast<float>(exp_avg_sq_q[index]) * old_variance_scale +
-        (1.0f - beta2) * gradient * gradient;
-    const float moment_q = nearbyintf(moment / new_moment_scale) + 128.0f;
-    const float variance_q = nearbyintf(variance / new_variance_scale);
-    exp_avg_q[index] = static_cast<uint8_t>(fminf(fmaxf(moment_q, 0.0f), 255.0f));
-    exp_avg_sq_q[index] = static_cast<uint8_t>(fminf(fmaxf(variance_q, 0.0f), 255.0f));
+    const float previous_moment = signed_codebook[exp_avg_q[index]] * old_moment_scale;
+    const float previous_variance = unsigned_codebook[exp_avg_sq_q[index]] * old_variance_scale;
+    const float moment = beta1 * previous_moment + (1.0f - beta1) * gradient;
+    const float variance = beta2 * previous_variance + (1.0f - beta2) * gradient * gradient;
+    float weight = load_model(model, index);
+    if (weight_decay != 0.0f) {
+      weight *= 1.0f - effective_lr * weight_decay;
+    }
+    const float denom = sqrtf(variance) / bias_correction2_sqrt + eps;
+    weight -= step_size * moment / denom;
+    const float normalized_moment = moment / fmaxf(new_moment_scale, 1.0e-30f);
+    const float normalized_variance = variance / fmaxf(new_variance_scale, 1.0e-30f);
+    exp_avg_q[index] = nearest_dynamic_code(normalized_moment, signed_codebook);
+    exp_avg_sq_q[index] = nearest_dynamic_code(normalized_variance, unsigned_codebook);
+    store_model(model, index, weight);
+  }
+}
+
+template <typename model_t, typename grad_t>
+__global__ void adamw_fp32_state_kernel(
+    model_t* model,
+    const grad_t* grad,
+    float* exp_avg,
+    float* exp_avg_sq,
+    int64_t numel,
+    float beta1,
+    float beta2,
+    float effective_lr,
+    float weight_decay,
+    float eps,
+    float step_size,
+    float bias_correction2_sqrt) {
+  for (int64_t index = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+       index < numel;
+       index += static_cast<int64_t>(blockDim.x) * gridDim.x) {
+    float moment = exp_avg[index];
+    float variance = exp_avg_sq[index];
+    const float weight = adamw_update(
+        load_model(model, index), load_grad(grad, index), moment, variance, beta1, beta2,
+        effective_lr, weight_decay, eps, step_size, bias_correction2_sqrt);
+    exp_avg[index] = moment;
+    exp_avg_sq[index] = variance;
+    store_model(model, index, weight);
   }
 }
 
@@ -495,6 +561,8 @@ void launch_adamw_8bit(
     torch::Tensor exp_avg_scale,
     torch::Tensor exp_avg_sq_q,
     torch::Tensor exp_avg_sq_scale,
+    torch::Tensor signed_codebook,
+    torch::Tensor unsigned_codebook,
     int64_t quant_block_size,
     float beta1,
     float beta2,
@@ -513,6 +581,8 @@ void launch_adamw_8bit(
       exp_avg_scale.data_ptr<float>(),
       exp_avg_sq_q.data_ptr<uint8_t>(),
       exp_avg_sq_scale.data_ptr<float>(),
+      signed_codebook.data_ptr<float>(),
+      unsigned_codebook.data_ptr<float>(),
       model.numel(),
       quant_block_size,
       beta1,
@@ -525,6 +595,28 @@ void launch_adamw_8bit(
   C10_CUDA_KERNEL_LAUNCH_CHECK();
 }
 
+template <typename model_t, typename grad_t>
+void launch_adamw_fp32_state(
+    torch::Tensor model,
+    torch::Tensor grad,
+    torch::Tensor exp_avg,
+    torch::Tensor exp_avg_sq,
+    float beta1,
+    float beta2,
+    float effective_lr,
+    float weight_decay,
+    float eps,
+    float step_size,
+    float bias_correction2_sqrt) {
+  constexpr int threads = 256;
+  const int blocks = static_cast<int>((model.numel() + threads - 1) / threads);
+  const auto stream = at::cuda::getCurrentCUDAStream();
+  adamw_fp32_state_kernel<model_t, grad_t><<<blocks, threads, 0, stream>>>(
+      model.data_ptr<model_t>(), grad.data_ptr<grad_t>(), exp_avg.data_ptr<float>(), exp_avg_sq.data_ptr<float>(),
+      model.numel(), beta1, beta2, effective_lr, weight_decay, eps, step_size, bias_correction2_sqrt);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
 }  // namespace
 
 void areno_adamw_fp32_master_step_cuda(
@@ -614,6 +706,8 @@ void areno_adamw_8bit_step_cuda(
     torch::Tensor exp_avg_scale,
     torch::Tensor exp_avg_sq_q,
     torch::Tensor exp_avg_sq_scale,
+    torch::Tensor signed_codebook,
+    torch::Tensor unsigned_codebook,
     int64_t quant_block_size,
     double beta1,
     double beta2,
@@ -625,11 +719,13 @@ void areno_adamw_8bit_step_cuda(
   c10::cuda::CUDAGuard guard(model.device());
   TORCH_CHECK(
       model.is_cuda() && grad.is_cuda() && exp_avg_q.is_cuda() && exp_avg_scale.is_cuda() &&
-          exp_avg_sq_q.is_cuda() && exp_avg_sq_scale.is_cuda(),
+          exp_avg_sq_q.is_cuda() && exp_avg_sq_scale.is_cuda() && signed_codebook.is_cuda() &&
+          unsigned_codebook.is_cuda(),
       "all 8-bit AdamW inputs must be CUDA tensors");
   TORCH_CHECK(
       model.is_contiguous() && grad.is_contiguous() && exp_avg_q.is_contiguous() && exp_avg_scale.is_contiguous() &&
-          exp_avg_sq_q.is_contiguous() && exp_avg_sq_scale.is_contiguous(),
+          exp_avg_sq_q.is_contiguous() && exp_avg_sq_scale.is_contiguous() && signed_codebook.is_contiguous() &&
+          unsigned_codebook.is_contiguous(),
       "all 8-bit AdamW inputs must be contiguous");
   TORCH_CHECK(model.numel() == grad.numel(), "model and gradient sizes must match");
   TORCH_CHECK(model.numel() == exp_avg_q.numel(), "model and first-moment sizes must match");
@@ -638,11 +734,14 @@ void areno_adamw_8bit_step_cuda(
   const int64_t block_count = (model.numel() + quant_block_size - 1) / quant_block_size;
   TORCH_CHECK(exp_avg_scale.numel() == block_count, "first-moment scale count must match quantization blocks");
   TORCH_CHECK(exp_avg_sq_scale.numel() == block_count, "second-moment scale count must match quantization blocks");
+  TORCH_CHECK(signed_codebook.numel() == 256, "signed dynamic codebook must have 256 entries");
+  TORCH_CHECK(unsigned_codebook.numel() == 256, "unsigned dynamic codebook must have 256 entries");
 
-#define LAUNCH_ADAMW8(MODEL_T, GRAD_T)                                                                \
-  launch_adamw_8bit<MODEL_T, GRAD_T>(                                                                 \
-      model, grad, exp_avg_q, exp_avg_scale, exp_avg_sq_q, exp_avg_sq_scale, quant_block_size, beta1, \
-      beta2, effective_lr, weight_decay, eps, step_size, bias_correction2_sqrt)
+#define LAUNCH_ADAMW8(MODEL_T, GRAD_T)                                                              \
+  launch_adamw_8bit<MODEL_T, GRAD_T>(                                                               \
+      model, grad, exp_avg_q, exp_avg_scale, exp_avg_sq_q, exp_avg_sq_scale, signed_codebook,        \
+      unsigned_codebook, quant_block_size, beta1, beta2, effective_lr, weight_decay, eps, step_size, \
+      bias_correction2_sqrt)
 
   if (model.scalar_type() == at::kBFloat16 && grad.scalar_type() == at::kBFloat16) {
     LAUNCH_ADAMW8(at::BFloat16, at::BFloat16);
@@ -657,3 +756,45 @@ void areno_adamw_8bit_step_cuda(
   }
 #undef LAUNCH_ADAMW8
 }
+
+void areno_adamw_fp32_state_step_cuda(
+    torch::Tensor model,
+    torch::Tensor grad,
+    torch::Tensor exp_avg,
+    torch::Tensor exp_avg_sq,
+    double beta1,
+    double beta2,
+    double effective_lr,
+    double weight_decay,
+    double eps,
+    double step_size,
+    double bias_correction2_sqrt) {
+  c10::cuda::CUDAGuard guard(model.device());
+  TORCH_CHECK(
+      model.is_cuda() && grad.is_cuda() && exp_avg.is_cuda() && exp_avg_sq.is_cuda(),
+      "all FP32-state AdamW inputs must be CUDA tensors");
+  TORCH_CHECK(
+      model.is_contiguous() && grad.is_contiguous() && exp_avg.is_contiguous() && exp_avg_sq.is_contiguous(),
+      "all FP32-state AdamW inputs must be contiguous");
+  TORCH_CHECK(model.numel() == grad.numel(), "model and gradient sizes must match");
+  TORCH_CHECK(model.numel() == exp_avg.numel(), "model and first-moment sizes must match");
+  TORCH_CHECK(model.numel() == exp_avg_sq.numel(), "model and second-moment sizes must match");
+
+#define LAUNCH_ADAMW_FP32_STATE(MODEL_T, GRAD_T)                                                     \
+  launch_adamw_fp32_state<MODEL_T, GRAD_T>(                                                         \
+      model, grad, exp_avg, exp_avg_sq, beta1, beta2, effective_lr, weight_decay, eps, step_size,    \
+      bias_correction2_sqrt)
+
+  if (model.scalar_type() == at::kBFloat16 && grad.scalar_type() == at::kBFloat16) {
+    LAUNCH_ADAMW_FP32_STATE(at::BFloat16, at::BFloat16);
+  } else if (model.scalar_type() == at::kBFloat16 && grad.scalar_type() == at::kFloat) {
+    LAUNCH_ADAMW_FP32_STATE(at::BFloat16, float);
+  } else if (model.scalar_type() == at::kFloat && grad.scalar_type() == at::kBFloat16) {
+    LAUNCH_ADAMW_FP32_STATE(float, at::BFloat16);
+  } else if (model.scalar_type() == at::kFloat && grad.scalar_type() == at::kFloat) {
+    LAUNCH_ADAMW_FP32_STATE(float, float);
+  } else {
+    TORCH_CHECK(false, "model and gradient must be bfloat16 or float32");
+  }
+#undef LAUNCH_ADAMW_FP32_STATE
+}
diff --git a/areno/accel/optimizer.py b/areno/accel/optimizer.py
index 0bba12ba..a339131b 100644
--- a/areno/accel/optimizer.py
+++ b/areno/accel/optimizer.py
@@ -76,6 +76,8 @@ def areno_adamw_8bit_step(
     exp_avg_scale: torch.Tensor,
     exp_avg_sq_q: torch.Tensor,
     exp_avg_sq_scale: torch.Tensor,
+    signed_codebook: torch.Tensor,
+    unsigned_codebook: torch.Tensor,
     *,
     block_size: int,
     beta1: float,
@@ -88,7 +90,16 @@ def areno_adamw_8bit_step(
 ) -> None:
     """Update block-quantized AdamW state without full FP32 moments."""
 
-    tensors = (model, grad, exp_avg_q, exp_avg_scale, exp_avg_sq_q, exp_avg_sq_scale)
+    tensors = (
+        model,
+        grad,
+        exp_avg_q,
+        exp_avg_scale,
+        exp_avg_sq_q,
+        exp_avg_sq_scale,
+        signed_codebook,
+        unsigned_codebook,
+    )
     if any(not tensor.is_cuda for tensor in tensors):
         raise ValueError("fused 8-bit AdamW requires CUDA tensors")
     if any(tensor.device != model.device for tensor in tensors[1:]):
@@ -101,6 +112,10 @@ def areno_adamw_8bit_step(
         raise TypeError("quantized Adam moments must use uint8")
     if exp_avg_scale.dtype != torch.float32 or exp_avg_sq_scale.dtype != torch.float32:
         raise TypeError("quantized Adam scales must use float32")
+    if signed_codebook.dtype != torch.float32 or unsigned_codebook.dtype != torch.float32:
+        raise TypeError("dynamic quantization codebooks must use float32")
+    if signed_codebook.numel() != 256 or unsigned_codebook.numel() != 256:
+        raise ValueError("dynamic quantization codebooks must contain 256 entries")
     if any(not tensor.is_contiguous() for tensor in tensors):
         raise ValueError("fused 8-bit AdamW requires contiguous tensors")
     if model.numel() != grad.numel() or model.numel() != exp_avg_q.numel():
@@ -119,6 +134,8 @@ def areno_adamw_8bit_step(
         exp_avg_scale,
         exp_avg_sq_q,
         exp_avg_sq_scale,
+        signed_codebook,
+        unsigned_codebook,
         block_size,
         beta1,
         beta2,
@@ -130,6 +147,54 @@ def areno_adamw_8bit_step(
     )
 
 
+@torch._dynamo.disable
+@torch.no_grad()
+def areno_adamw_fp32_state_step(
+    model: torch.Tensor,
+    grad: torch.Tensor,
+    exp_avg: torch.Tensor,
+    exp_avg_sq: torch.Tensor,
+    *,
+    beta1: float,
+    beta2: float,
+    effective_lr: float,
+    weight_decay: float,
+    eps: float,
+    step_size: float,
+    bias_correction2_sqrt: float,
+) -> None:
+    """Update BF16/FP32 weights with persistent FP32 Adam moments."""
+
+    tensors = (model, grad, exp_avg, exp_avg_sq)
+    if any(not tensor.is_cuda for tensor in tensors):
+        raise ValueError("fused FP32-state AdamW requires CUDA tensors")
+    if any(tensor.device != model.device for tensor in tensors[1:]):
+        raise ValueError("fused FP32-state AdamW requires every tensor on the model device")
+    if model.dtype not in {torch.bfloat16, torch.float32}:
+        raise TypeError(f"fused FP32-state AdamW requires bfloat16 or float32 model weights, got {model.dtype}")
+    if grad.dtype not in {torch.bfloat16, torch.float32}:
+        raise TypeError(f"fused FP32-state AdamW requires bfloat16 or float32 gradients, got {grad.dtype}")
+    if exp_avg.dtype != torch.float32 or exp_avg_sq.dtype != torch.float32:
+        raise TypeError("FP32-state Adam moments must use float32")
+    if any(not tensor.is_contiguous() for tensor in tensors):
+        raise ValueError("fused FP32-state AdamW requires contiguous tensors")
+    if any(tensor.numel() != model.numel() for tensor in tensors[1:]):
+        raise ValueError("model, gradient, and FP32 moments must have the same number of elements")
+    extension().areno_adamw_fp32_state_step(
+        model,
+        grad,
+        exp_avg,
+        exp_avg_sq,
+        beta1,
+        beta2,
+        effective_lr,
+        weight_decay,
+        eps,
+        step_size,
+        bias_correction2_sqrt,
+    )
+
+
 @torch._dynamo.disable
 @torch.no_grad()
 def areno_adamw_4bit_step(
@@ -200,4 +265,9 @@ def areno_adamw_4bit_step(
     )
 
 
-__all__ = ["areno_adamw_4bit_step", "areno_adamw_8bit_step", "areno_adamw_fp32_master_step"]
+__all__ = [
+    "areno_adamw_4bit_step",
+    "areno_adamw_8bit_step",
+    "areno_adamw_fp32_master_step",
+    "areno_adamw_fp32_state_step",
+]
diff --git a/areno/api/backend/mlx/backend.py b/areno/api/backend/mlx/backend.py
index c76c70bf..5175b083 100644
--- a/areno/api/backend/mlx/backend.py
+++ b/areno/api/backend/mlx/backend.py
@@ -92,7 +92,10 @@ def initialize(self, ctx: Context):
         self._validate_tokenizer(ctx.tokenizer)
         optimizer_config = self.config.optimizer
         self.provider.configure_trainability(optimizer_config)
-        self.optimizer, self._optimizer_groups = build_optimizer(optimizer_config)
+        self.optimizer, self._optimizer_groups = build_optimizer(
+            optimizer_config,
+            state_precision_for_parameter=self.provider.optimizer_state_precision,
+        )
         if self.config.gradient_checkpointing:
             self._enable_gradient_checkpointing()
         self.model.train()
diff --git a/areno/api/backend/mlx/optimizer.py b/areno/api/backend/mlx/optimizer.py
index fcdf1d64..320a52da 100644
--- a/areno/api/backend/mlx/optimizer.py
+++ b/areno/api/backend/mlx/optimizer.py
@@ -2,12 +2,20 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from typing import Any
 
 from areno.api.backend.mlx.provider import parameter_group
+from areno.engine.optim.dynamic_quant import SIGNED_DYNAMIC_MAP, UNSIGNED_DYNAMIC_MAP
 
+_MLX_CODEBOOK_CACHE: dict[bool, Any] = {}
 
-def build_optimizer(config: dict[str, Any]):
+
+def build_optimizer(
+    config: dict[str, Any],
+    *,
+    state_precision_for_parameter: Callable[[str, Any], str] | None = None,
+):
     """Build AdamW groups matching CUDA policy/tower/projector controls."""
 
     import mlx.optimizers as optim
@@ -16,14 +24,14 @@ def build_optimizer(config: dict[str, Any]):
     filters = []
     if config.get("unfreeze_multimodal_tower"):
         tower_config = _group_config(config, "tower")
-        groups.append(("tower", _adamw(tower_config), tower_config))
+        groups.append(("tower", _adamw(tower_config, state_precision_for_parameter), tower_config))
         filters.append(lambda path, _: parameter_group(path) == "tower")
     if config.get("unfreeze_multimodal_projector"):
         projector_config = _group_config(config, "projector")
-        groups.append(("projector", _adamw(projector_config), projector_config))
+        groups.append(("projector", _adamw(projector_config, state_precision_for_parameter), projector_config))
         filters.append(lambda path, _: parameter_group(path) == "projector")
     model_config = dict(config)
-    groups.append(("model", _adamw(model_config), model_config))
+    groups.append(("model", _adamw(model_config, state_precision_for_parameter), model_config))
     if len(groups) == 1:
         return groups[0][1], groups
     optimizer_type = _streaming_multi_optimizer_class() if config.get("adam_8bit") else optim.MultiOptimizer
@@ -78,7 +86,10 @@ def _group_config(config: dict[str, Any], group: str) -> dict[str, Any]:
     return result
 
 
-def _adamw(config: dict[str, Any]):
+def _adamw(
+    config: dict[str, Any],
+    state_precision_for_parameter: Callable[[str, Any], str] | None = None,
+):
     import mlx.optimizers as optim
 
     kwargs = {
@@ -88,7 +99,10 @@ def _adamw(config: dict[str, Any]):
     }
     if not config.get("adam_8bit"):
         return optim.AdamW(**kwargs, bias_correction=True)
-    return _quantized_adamw_class()(**kwargs)
+    return _quantized_adamw_class()(
+        **kwargs,
+        state_precision_for_parameter=state_precision_for_parameter,
+    )
 
 
 def _quantized_adamw_class():
@@ -96,7 +110,7 @@ def _quantized_adamw_class():
     from mlx.optimizers import Optimizer
 
     class AdamW8bit(Optimizer):
-        """AdamW with blockwise uint8 first-moment and root-second-moment storage."""
+        """Paper-compatible blockwise dynamic AdamW with FP32 embedding states."""
 
         def __init__(
             self,
@@ -104,8 +118,9 @@ def __init__(
             betas=(0.9, 0.999),
             eps: float = 1e-8,
             weight_decay: float = 0.01,
-            block_size: int = 256,
+            block_size: int = 128,
             update_blocks: int = 8192,
+            state_precision_for_parameter: Callable[[str, Any], str] | None = None,
         ) -> None:
             super().__init__()
             self._maybe_schedule("learning_rate", learning_rate)
@@ -114,6 +129,7 @@ def __init__(
             self.weight_decay = float(weight_decay)
             self.block_size = int(block_size)
             self.update_blocks = int(update_blocks)
+            self.state_precision_for_parameter = state_precision_for_parameter
 
         def init_single(self, parameter, state: dict) -> None:
             size = int(parameter.size)
@@ -128,6 +144,28 @@ def apply_single(self, gradient, parameter, state: dict):
             bias_correction2_sqrt = mx.sqrt(1.0 - beta2**step)
             size = int(state["size"])
             initialized = bool(state["initialized"])
+            precision = str(state.get("precision", "8bit"))
+            if precision == "fp32":
+                grad = gradient.astype(mx.float32)
+                values = parameter.astype(mx.float32)
+                if initialized:
+                    m = state["m"]
+                    v = state["v"]
+                else:
+                    m = mx.zeros_like(values, dtype=mx.float32)
+                    v = mx.zeros_like(values, dtype=mx.float32)
+                m = beta1 * m + (1.0 - beta1) * grad
+                v = beta2 * v + (1.0 - beta2) * mx.square(grad)
+                denom = mx.sqrt(v) / bias_correction2_sqrt + self.eps
+                updated = (values * (1.0 - lr * self.weight_decay) - (lr / bias_correction1) * m / denom).astype(
+                    parameter.dtype
+                )
+                state["m"] = m
+                state["v"] = v
+                state["initialized"] = True
+                mx.eval(updated, m, v)
+                mx.clear_cache()
+                return updated
             block_count = (size + self.block_size - 1) // self.block_size
             grad = gradient.reshape(-1)
             values = parameter.reshape(-1)
@@ -154,19 +192,18 @@ def apply_single(self, gradient, parameter, state: dict):
                         state["m_scale"][block_start:block_end],
                         signed=True,
                     )
-                    v_root = _dequant_blocks(
+                    v = _dequant_blocks(
                         state["v_q"][value_start:padded_end],
                         state["v_scale"][block_start:block_end],
                         signed=False,
                     )
-                    v = mx.square(v_root)
                     m = beta1 * m + (1.0 - beta1) * grad_chunk
                     v = beta2 * v + (1.0 - beta2) * mx.square(grad_chunk)
                 else:
                     m = (1.0 - beta1) * grad_chunk
                     v = (1.0 - beta2) * mx.square(grad_chunk)
                 next_m_q, next_m_scale = _quantize_signed(m, self.block_size)
-                next_v_q, next_v_scale = _quantize_unsigned(mx.sqrt(v), self.block_size)
+                next_v_q, next_v_scale = _quantize_unsigned(v, self.block_size)
                 denom = mx.sqrt(v) / bias_correction2_sqrt + self.eps
                 updated = (value_chunk * (1.0 - lr * self.weight_decay) - (lr / bias_correction1) * m / denom)[
                     :actual
@@ -194,6 +231,18 @@ def update_streaming(self, model, gradients: dict) -> None:
             self._begin_streaming_step()
             _apply_streaming_leaves(model, gradients, lambda _: self)
 
+        def prepare_parameter_state(self, path: str, parameter: Any, state: dict) -> None:
+            if "precision" in state:
+                return
+            precision = (
+                "8bit"
+                if self.state_precision_for_parameter is None
+                else str(self.state_precision_for_parameter(path, parameter))
+            )
+            if precision not in {"8bit", "fp32"}:
+                raise ValueError(f"unsupported MLX AdamW8bit state precision: {precision!r}")
+            state["precision"] = precision
+
         def _begin_streaming_step(self) -> None:
             for name, scheduler in self._schedulers.items():
                 self.state[name] = scheduler(self.step)
@@ -235,6 +284,9 @@ def _apply_streaming_leaves(model: Any, gradients: dict, optimizer_for_path) ->
         optimizer = optimizer_for_path(path)
         state = _tree_get(optimizer.state, path)
         parameter = _model_parameter(model, path)
+        prepare = getattr(optimizer, "prepare_parameter_state", None)
+        if prepare is not None:
+            prepare(path, parameter, state)
         updated = optimizer.apply_single(gradient, parameter, state)
         _set_model_parameter(model, path, updated)
         _tree_set(gradients, path, None)
@@ -245,7 +297,7 @@ def _apply_streaming_leaves(model: Any, gradients: dict, optimizer_for_path) ->
 def _tree_get(tree: Any, path: str) -> Any:
     current = tree
     for part in path.split("."):
-        current = current[int(part)] if isinstance(current, (list, tuple)) else current[part]
+        current = current[int(part)] if isinstance(current, list | tuple) else current[part]
     return current
 
 
@@ -253,7 +305,7 @@ def _tree_set(tree: Any, path: str, value: Any) -> None:
     parts = path.split(".")
     current = tree
     for part in parts[:-1]:
-        current = current[int(part)] if isinstance(current, (list, tuple)) else current[part]
+        current = current[int(part)] if isinstance(current, list | tuple) else current[part]
     final = parts[-1]
     if isinstance(current, list):
         current[int(final)] = value
@@ -264,7 +316,7 @@ def _tree_set(tree: Any, path: str, value: Any) -> None:
 def _model_parameter(model: Any, path: str):
     current = model
     for part in path.split("."):
-        current = current[int(part)] if isinstance(current, (list, tuple)) else getattr(current, part)
+        current = current[int(part)] if isinstance(current, list | tuple) else getattr(current, part)
     return current
 
 
@@ -272,7 +324,7 @@ def _set_model_parameter(model: Any, path: str, value: Any) -> None:
     parts = path.split(".")
     current = model
     for part in parts[:-1]:
-        current = current[int(part)] if isinstance(current, (list, tuple)) else getattr(current, part)
+        current = current[int(part)] if isinstance(current, list | tuple) else getattr(current, part)
     final = parts[-1]
     if isinstance(current, list):
         current[int(final)] = value
@@ -291,28 +343,44 @@ def _blocked(value, block_size: int):
 
 
 def _quantize_signed(value, block_size: int):
-    import mlx.core as mx
-
-    blocks = _blocked(value, block_size)
-    scale = mx.maximum(mx.max(mx.abs(blocks), axis=1, keepdims=True) / 127.0, mx.array(1e-12))
-    quantized = mx.clip(mx.round(blocks / scale), -127, 127).astype(mx.int16) + 128
-    return quantized.astype(mx.uint8).reshape(-1), scale
+    return _quantize_dynamic(value, block_size, signed=True)
 
 
 def _quantize_unsigned(value, block_size: int):
+    return _quantize_dynamic(value, block_size, signed=False)
+
+
+def _quantize_dynamic(value, block_size: int, *, signed: bool):
     import mlx.core as mx
 
     blocks = _blocked(value, block_size)
-    scale = mx.maximum(mx.max(blocks, axis=1, keepdims=True) / 255.0, mx.array(1e-12))
-    quantized = mx.clip(mx.round(blocks / scale), 0, 255).astype(mx.uint8)
-    return quantized.reshape(-1), scale
+    scale = mx.max(mx.abs(blocks) if signed else mx.maximum(blocks, 0.0), axis=1, keepdims=True)
+    normalized = blocks / mx.maximum(scale, mx.array(1.0e-30, dtype=mx.float32))
+    if not signed:
+        normalized = mx.maximum(normalized, 0.0)
+    codebook = _mlx_dynamic_codebook(signed=signed)
+    boundaries = (codebook[:-1] + codebook[1:]) * 0.5
+    quantized = mx.searchsorted(boundaries, normalized).astype(mx.uint8)
+    return quantized.reshape(-1), scale.astype(mx.float32)
 
 
 def _dequant_blocks(quantized, scale, *, signed: bool):
-    values = quantized.reshape(scale.shape[0], -1).astype(scale.dtype)
-    if signed:
-        values = values - 128.0
+    import mlx.core as mx
+
+    codebook = _mlx_dynamic_codebook(signed=signed)
+    values = codebook[quantized.reshape(scale.shape[0], -1).astype(mx.uint32)]
     return (values * scale).reshape(-1)
 
 
+def _mlx_dynamic_codebook(*, signed: bool):
+    import mlx.core as mx
+
+    codebook = _MLX_CODEBOOK_CACHE.get(signed)
+    if codebook is None:
+        values = SIGNED_DYNAMIC_MAP if signed else UNSIGNED_DYNAMIC_MAP
+        codebook = mx.array(values, dtype=mx.float32)
+        _MLX_CODEBOOK_CACHE[signed] = codebook
+    return codebook
+
+
 __all__ = ["apply_optimizer_update", "build_optimizer", "materialize_optimizer_update", "set_group_learning_rates"]
diff --git a/areno/api/backend/mlx/provider.py b/areno/api/backend/mlx/provider.py
index ffa275a0..c9421ea4 100644
--- a/areno/api/backend/mlx/provider.py
+++ b/areno/api/backend/mlx/provider.py
@@ -65,6 +65,43 @@ def generation_model(self):
     def configure_trainability(self, optimizer_config: dict[str, Any]) -> None:
         del optimizer_config
 
+    def optimizer_state_precision(self, path: str, parameter: Any) -> str:
+        """Return role-aware optimizer-state precision for one MLX parameter."""
+
+        del path
+        for embedding in self._token_embedding_modules():
+            if getattr(embedding, "weight", None) is parameter:
+                return "fp32"
+        return "8bit"
+
+    def _token_embedding_modules(self) -> tuple[Any, ...]:
+        language_model = self.generation_model
+        body = getattr(language_model, "model", None)
+        modules = []
+        seen: set[int] = set()
+
+        def add(module: Any) -> None:
+            if module is None:
+                return
+            if isinstance(module, dict):
+                for child in module.values():
+                    add(child)
+                return
+            if isinstance(module, list | tuple):
+                for child in module:
+                    add(child)
+                return
+            if id(module) not in seen:
+                modules.append(module)
+                seen.add(id(module))
+
+        for owner in (body, language_model):
+            if owner is None:
+                continue
+            for name in ("embed_tokens", "word_embeddings", "embed_tokens_per_layer"):
+                add(getattr(owner, name, None))
+        return tuple(modules)
+
     def prepare_generation_prompt(self, tokens: list[int], features: dict[str, Any] | None) -> dict[str, Any]:
         if features is not None:
             raise ValueError("text-only MLX checkpoints cannot consume multimodal prompt features")
diff --git a/areno/api/data.py b/areno/api/data.py
index d3de09f8..a13c5b1a 100644
--- a/areno/api/data.py
+++ b/areno/api/data.py
@@ -8,8 +8,19 @@
 
 from __future__ import annotations
 
+import hashlib
+import json
+import math
+import random
+from collections import Counter
+from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Literal
+
+DATASET_MIX_METADATA_KEY = "__areno_meta__"
+DatasetExhaustionPolicy = Literal["stop", "cycle", "renormalize"]
+DATASET_MIX_SAMPLER_VERSION = 1
+DATASET_MIX_WEIGHT_UNIT = "sample"
 
 
 @dataclass(slots=True)
@@ -48,3 +59,271 @@ def prompts(self) -> list[str]:
         """Return raw prompt strings in batch order for rollout."""
 
         return [item.prompt for item in self.items]
+
+
+@dataclass(frozen=True, slots=True)
+class DatasetMixSource:
+    """One named, weighted map-style dataset used by ``WeightedMixedDataset``."""
+
+    name: str
+    dataset: Sequence
+    weight: float
+
+
+@dataclass(frozen=True, slots=True)
+class _DatasetMixEntry:
+    source_index: int
+    row_index: int
+    cycle: int
+
+
+class WeightedMixedDataset:
+    """Deterministically interleave weighted map-style datasets.
+
+    ``stop`` ends when a selected source is exhausted. ``cycle`` restarts
+    exhausted sources until ``samples_per_epoch`` is reached, or until every
+    source has exhausted once when no budget is supplied. ``renormalize``
+    removes exhausted sources and continues with the remaining weights,
+    emitting every source row exactly once.
+    """
+
+    def __init__(
+        self,
+        sources: Sequence[DatasetMixSource],
+        *,
+        seed: int,
+        exhaustion: DatasetExhaustionPolicy,
+        shuffle_within_sources: bool = True,
+        samples_per_epoch: int | None = None,
+    ) -> None:
+        self.sources = tuple(sources)
+        self.seed = seed
+        self.exhaustion = exhaustion
+        self.shuffle_within_sources = shuffle_within_sources
+        self.samples_per_epoch = samples_per_epoch
+        self.epoch = 0
+        self._normalized_weights: tuple[float, ...] = ()
+        self._validate()
+        self._entries: list[_DatasetMixEntry] = []
+        self._termination_reason = ""
+        self._summary_cache: dict[str, Any] = {}
+        self.set_epoch(0)
+
+    def _validate(self) -> None:
+        if len(self.sources) < 2:
+            raise ValueError("dataset mix requires at least two sources")
+        if isinstance(self.seed, bool) or not isinstance(self.seed, int) or not 0 <= self.seed < 2**63:
+            raise ValueError("dataset mix seed must be an integer in [0, 2^63)")
+        if self.exhaustion not in {"stop", "cycle", "renormalize"}:
+            raise ValueError("dataset mix exhaustion must be one of: stop, cycle, renormalize")
+        if self.samples_per_epoch is not None and (
+            isinstance(self.samples_per_epoch, bool)
+            or not isinstance(self.samples_per_epoch, int)
+            or self.samples_per_epoch <= 0
+        ):
+            raise ValueError("dataset mix samples_per_epoch must be a positive integer")
+        if self.samples_per_epoch is not None and self.exhaustion != "cycle":
+            raise ValueError("dataset mix samples_per_epoch is only supported with exhaustion='cycle'")
+        if not isinstance(self.shuffle_within_sources, bool):
+            raise ValueError("dataset mix shuffle_within_sources must be a boolean")
+
+        names: set[str] = set()
+        numeric_weights: list[float] = []
+        for source in self.sources:
+            if not isinstance(source.name, str) or not source.name.strip():
+                raise ValueError("dataset mix source name must be a non-empty string")
+            if source.name != source.name.strip():
+                raise ValueError("dataset mix source name must not have surrounding whitespace")
+            if not source.name.isprintable():
+                raise ValueError(f"dataset mix source name contains non-printable characters: {source.name!r}")
+            if source.name in names:
+                raise ValueError(f"duplicate dataset mix source name: {source.name}")
+            names.add(source.name)
+            try:
+                numeric_weight = float(source.weight)
+            except (OverflowError, TypeError, ValueError):
+                numeric_weight = math.nan
+            if isinstance(source.weight, bool) or not math.isfinite(numeric_weight) or numeric_weight <= 0:
+                raise ValueError(f"dataset mix source '{source.name}' weight must be finite and positive")
+            numeric_weights.append(numeric_weight)
+            if not hasattr(source.dataset, "__len__") or not hasattr(source.dataset, "__getitem__"):
+                raise ValueError(f"dataset mix source '{source.name}' must support len() and indexed access")
+            if len(source.dataset) == 0:
+                raise ValueError(f"dataset mix source '{source.name}' is empty")
+            first = source.dataset[0]
+            if not isinstance(first, Mapping):
+                raise ValueError(f"dataset mix source '{source.name}' rows must be mappings")
+            if DATASET_MIX_METADATA_KEY in first:
+                raise ValueError(
+                    f"dataset mix source '{source.name}' contains reserved field '{DATASET_MIX_METADATA_KEY}'"
+                )
+        max_weight = max(numeric_weights)
+        scaled_weights = [weight / max_weight for weight in numeric_weights]
+        if any(weight == 0.0 for weight in scaled_weights):
+            raise ValueError("dataset mix weights have an unsupported numeric range")
+        scaled_total = sum(scaled_weights)
+        normalized_weights = tuple(weight / scaled_total for weight in scaled_weights)
+        if any(weight == 0.0 for weight in normalized_weights):
+            raise ValueError("dataset mix weights have an unsupported numeric range")
+        self._normalized_weights = normalized_weights
+
+    def set_epoch(self, epoch: int) -> None:
+        """Build the deterministic schedule for one epoch."""
+
+        if isinstance(epoch, bool) or not isinstance(epoch, int) or epoch < 0:
+            raise ValueError("dataset mix epoch must be a non-negative integer")
+        if epoch == self.epoch and self._entries:
+            return
+        self.epoch = epoch
+        self._entries, self._termination_reason = self._build_schedule()
+        self._summary_cache = self._build_summary()
+
+    def _build_schedule(self) -> tuple[list[_DatasetMixEntry], str]:
+        rng = random.Random(_stable_seed(self.seed, self.epoch, "source-selection"))
+        orders = [self._source_order(index, cycle=0) for index in range(len(self.sources))]
+        positions = [0] * len(self.sources)
+        cycles = [0] * len(self.sources)
+        exhausted_once: set[int] = set()
+        active = list(range(len(self.sources)))
+        entries: list[_DatasetMixEntry] = []
+
+        while active:
+            if self.samples_per_epoch is not None and len(entries) >= self.samples_per_epoch:
+                return entries, "samples_per_epoch"
+            source_index = _weighted_choice(rng, active, [self._normalized_weights[index] for index in active])
+
+            row_index = orders[source_index][positions[source_index]]
+            positions[source_index] += 1
+            entries.append(_DatasetMixEntry(source_index=source_index, row_index=row_index, cycle=cycles[source_index]))
+            if positions[source_index] < len(orders[source_index]):
+                continue
+
+            exhausted_once.add(source_index)
+            if self.exhaustion == "stop":
+                return entries, f"source_exhausted:{self.sources[source_index].name}"
+            if self.exhaustion == "renormalize":
+                active.remove(source_index)
+                continue
+            if self.samples_per_epoch is None and len(exhausted_once) == len(self.sources):
+                return entries, "all_sources_exhausted_once"
+            cycles[source_index] += 1
+            orders[source_index] = self._source_order(source_index, cycle=cycles[source_index])
+            positions[source_index] = 0
+
+        return entries, "all_sources_exhausted"
+
+    def _source_order(self, source_index: int, *, cycle: int) -> Sequence[int]:
+        if not self.shuffle_within_sources:
+            return range(len(self.sources[source_index].dataset))
+        order = list(range(len(self.sources[source_index].dataset)))
+        source = self.sources[source_index]
+        random.Random(_stable_seed(self.seed, self.epoch, source.name, cycle)).shuffle(order)
+        return order
+
+    def __len__(self) -> int:
+        return len(self._entries)
+
+    def __getitem__(self, index: int) -> dict[str, Any]:
+        entry = self._entries[index]
+        source = self.sources[entry.source_index]
+        record = dict(source.dataset[entry.row_index])
+        if DATASET_MIX_METADATA_KEY in record:
+            raise ValueError(
+                f"dataset mix source '{source.name}' row {entry.row_index} contains reserved field "
+                f"'{DATASET_MIX_METADATA_KEY}'"
+            )
+        record[DATASET_MIX_METADATA_KEY] = {
+            "source": source.name,
+            "source_index": entry.row_index,
+            "cycle": entry.cycle,
+        }
+        return record
+
+    def summary(self) -> dict[str, Any]:
+        """Return sample-free metadata suitable for logs and JSON artifacts."""
+
+        return {
+            **self._summary_cache,
+            "warnings": list(self._summary_cache["warnings"]),
+            "sources": [dict(source) for source in self._summary_cache["sources"]],
+        }
+
+    def _build_summary(self) -> dict[str, Any]:
+        selected = Counter(entry.source_index for entry in self._entries)
+        duplicates = Counter(entry.source_index for entry in self._entries if entry.cycle > 0)
+        total = len(self._entries)
+        source_summaries = []
+        warnings = []
+        for index, source in enumerate(self.sources):
+            count = selected[index]
+            expected_rows = (
+                self.samples_per_epoch * self._normalized_weights[index] if self.samples_per_epoch is not None else None
+            )
+            if expected_rows is not None and expected_rows < 1:
+                warnings.append(
+                    f"source '{source.name}' has expected_rows={expected_rows:.6g}; "
+                    "it may receive zero samples in an epoch"
+                )
+            source_summaries.append(
+                {
+                    "name": source.name,
+                    "weight_requested": self._normalized_weights[index],
+                    "expected_rows": expected_rows,
+                    "rows_available": len(source.dataset),
+                    "rows_selected": count,
+                    "duplicates": duplicates[index],
+                    "observed_proportion": count / total if total else 0.0,
+                }
+            )
+        schedule_hash = hashlib.sha256()
+        for entry in self._entries:
+            schedule_hash.update(f"{self.sources[entry.source_index].name}:{entry.row_index}:{entry.cycle}\n".encode())
+        mix_spec = {
+            "sampler_version": DATASET_MIX_SAMPLER_VERSION,
+            "weight_unit": DATASET_MIX_WEIGHT_UNIT,
+            "seed": self.seed,
+            "policy": self.exhaustion,
+            "shuffle_within_sources": self.shuffle_within_sources,
+            "samples_per_epoch": self.samples_per_epoch,
+            "sources": [
+                {
+                    "name": source.name,
+                    "rows_available": len(source.dataset),
+                    "weight": self._normalized_weights[index],
+                }
+                for index, source in enumerate(self.sources)
+            ],
+        }
+        serialized_mix_spec = json.dumps(mix_spec, sort_keys=True, separators=(",", ":"), ensure_ascii=False)
+        mix_spec_hash = hashlib.sha256(serialized_mix_spec.encode()).hexdigest()
+        return {
+            "version": 1,
+            "sampler_version": DATASET_MIX_SAMPLER_VERSION,
+            "weight_unit": DATASET_MIX_WEIGHT_UNIT,
+            "seed": self.seed,
+            "epoch": self.epoch,
+            "policy": self.exhaustion,
+            "shuffle_within_sources": self.shuffle_within_sources,
+            "samples_per_epoch": self.samples_per_epoch,
+            "planned_rows": total,
+            "termination_reason": self._termination_reason,
+            "mix_spec_hash": f"sha256:{mix_spec_hash}",
+            "schedule_hash": f"sha256:{schedule_hash.hexdigest()}",
+            "warnings": warnings,
+            "sources": source_summaries,
+        }
+
+
+def _weighted_choice(rng: random.Random, candidates: Sequence[int], weights: Sequence[float]) -> int:
+    threshold = rng.random() * sum(weights)
+    cumulative = 0.0
+    for candidate, weight in zip(candidates, weights, strict=True):
+        cumulative += weight
+        if threshold < cumulative:
+            return candidate
+    return candidates[-1]
+
+
+def _stable_seed(*parts: object) -> int:
+    payload = "\0".join(str(part) for part in parts).encode()
+    return int.from_bytes(hashlib.sha256(payload).digest()[:8], "big")
diff --git a/areno/api/dataset_mix_artifacts.py b/areno/api/dataset_mix_artifacts.py
new file mode 100644
index 00000000..4e36a588
--- /dev/null
+++ b/areno/api/dataset_mix_artifacts.py
@@ -0,0 +1,35 @@
+"""Structured, sample-free artifacts for dataset-mixing plans."""
+
+from __future__ import annotations
+
+import json
+import os
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any
+
+
+def write_dataset_mix_plan(
+    summary: Mapping[str, Any],
+    metrics_log_dir: str | None,
+    *,
+    process_id: int | None = None,
+) -> Path | None:
+    """Write one deterministic plan artifact for the summary's epoch."""
+
+    if not metrics_log_dir:
+        return None
+
+    epoch = summary.get("epoch")
+    if isinstance(epoch, bool) or not isinstance(epoch, int) or epoch < 0:
+        raise ValueError("dataset mix summary epoch must be a non-negative integer")
+
+    path = Path(metrics_log_dir)
+    path.mkdir(parents=True, exist_ok=True)
+    pid = os.getpid() if process_id is None else process_id
+    artifact_path = path / f"dataset_mix_plan.{pid}.epoch-{epoch}.json"
+    artifact_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
+    return artifact_path
+
+
+__all__ = ["write_dataset_mix_plan"]
diff --git a/areno/api/trainer_config.py b/areno/api/trainer_config.py
index 7db26af3..ba37dede 100644
--- a/areno/api/trainer_config.py
+++ b/areno/api/trainer_config.py
@@ -28,11 +28,16 @@ class TrainerConfig:
 
     algo: str
     ckpt: str
-    dataset_path: str
+    dataset_path: str | None
     backend: str | None = None
     base_model_name_or_path: str | None = field(default=None, kw_only=True)
     model_hub: str = "modelscope"
     dataset_loader_fn: str | None = None
+    dataset_mix_config: str | None = field(default=None, kw_only=True)
+    dataset_sources: tuple[str, ...] = field(default=(), kw_only=True)
+    dataset_mix_seed: int = field(default=42, kw_only=True)
+    dataset_mix_exhaustion: str = field(default="cycle", kw_only=True)
+    dataset_mix_samples_per_epoch: int | None = field(default=None, kw_only=True)
     save_path: str | None = None
     save_interval: int = 100
     epochs: int = 10
@@ -83,6 +88,41 @@ class TrainerConfig:
     reference_mode: Literal["independent", "reuse_actor_base"] = "independent"
 
     def __post_init__(self) -> None:
+        dataset_inputs = sum(
+            (
+                self.dataset_path is not None,
+                self.dataset_mix_config is not None,
+                bool(self.dataset_sources),
+            )
+        )
+        if dataset_inputs == 0:
+            raise ValueError("one of dataset_path, dataset_mix_config, or dataset_sources is required")
+        if dataset_inputs > 1:
+            raise ValueError("dataset_path, dataset_mix_config, and dataset_sources are mutually exclusive")
+        if not self.dataset_sources and (
+            self.dataset_mix_seed != 42
+            or self.dataset_mix_exhaustion != "cycle"
+            or self.dataset_mix_samples_per_epoch is not None
+        ):
+            raise ValueError("inline dataset mix settings apply only when dataset_sources is configured")
+        if self.dataset_sources and len(self.dataset_sources) < 2:
+            raise ValueError("dataset_sources must contain at least two entries")
+        if (self.dataset_mix_config is not None or self.dataset_sources) and self.algo != "sft":
+            raise ValueError("dataset mixing currently supports algo='sft' only")
+        if isinstance(self.dataset_mix_seed, bool) or not isinstance(self.dataset_mix_seed, int):
+            raise ValueError("dataset_mix_seed must be an integer")
+        if not 0 <= self.dataset_mix_seed < 2**63:
+            raise ValueError("dataset_mix_seed must be in [0, 2^63)")
+        if self.dataset_mix_exhaustion not in {"stop", "cycle", "renormalize"}:
+            raise ValueError("dataset_mix_exhaustion must be one of: stop, cycle, renormalize")
+        if self.dataset_mix_samples_per_epoch is not None and (
+            isinstance(self.dataset_mix_samples_per_epoch, bool)
+            or not isinstance(self.dataset_mix_samples_per_epoch, int)
+            or self.dataset_mix_samples_per_epoch <= 0
+        ):
+            raise ValueError("dataset_mix_samples_per_epoch must be a positive integer")
+        if self.dataset_mix_samples_per_epoch is not None and self.dataset_mix_exhaustion != "cycle":
+            raise ValueError("dataset_mix_samples_per_epoch requires dataset_mix_exhaustion='cycle'")
         if self.backend is None:
             from areno.api.config import default_backend_type
 
diff --git a/areno/api/trainers/sft.py b/areno/api/trainers/sft.py
index 8774b297..6cfe24bb 100644
--- a/areno/api/trainers/sft.py
+++ b/areno/api/trainers/sft.py
@@ -18,12 +18,16 @@
 
 import logging
 import time
+from collections import Counter
+from collections.abc import Mapping
 from pathlib import Path
 from typing import Any
 
 import areno.api
 from areno.api.dashboard import record_dashboard_state
+from areno.api.data import DATASET_MIX_METADATA_KEY
 from areno.api.data_utils import prompt_response_to_tokens_and_mask
+from areno.api.dataset_mix_artifacts import write_dataset_mix_plan
 from areno.api.multimodal import (
     encode_multimodal_prompt,
     expand_image_tokens,
@@ -65,13 +69,37 @@ def _fit_initialized(self) -> None:
         configure_chat_template_enable_thinking(processor, getattr(self.config, "chat_template_enable_thinking", None))
         step = 0
         for epoch in range(self.config.epochs):
+            set_epoch = getattr(self.dataset, "set_epoch", None)
+            if callable(set_epoch):
+                set_epoch(epoch)
+            mix_summary = getattr(self.dataset, "summary", None)
+            if callable(mix_summary):
+                resolved_mix_summary = mix_summary()
+                if "schedule_hash" in resolved_mix_summary:
+                    write_dataset_mix_plan(
+                        resolved_mix_summary,
+                        getattr(self.config, "metrics_log_dir", None),
+                    )
+                self.logger.info("epoch=%d stage=dataset_mix_plan dataset_mix=%s", epoch, resolved_mix_summary)
+                mix_source_names = [
+                    source["name"] for source in resolved_mix_summary.get("sources", []) if "name" in source
+                ]
+            else:
+                mix_source_names = []
+            mix_progress = {
+                "scheduled": Counter(),
+                "filtered": Counter(),
+                "trained": Counter(),
+                "target_tokens": Counter(),
+            }
             self.logger.info("epoch=%d stage=epoch_start", epoch)
             record_dashboard_state(self.areno, stage="epoch_start", epoch=epoch, step=step, role="policy")
-            for train_batch in self._iter_train_batches(
+            for train_batch, batch_mix_counts, batch_mix_target_tokens in self._iter_train_batches(
                 tokenizer,
                 processor,
                 max_prompt_tokens=self.config.max_prompt_tokens,
                 max_new_tokens=self.config.max_new_tokens,
+                mix_progress=mix_progress,
             ):
                 if not train_batch:
                     continue
@@ -95,27 +123,56 @@ def _fit_initialized(self) -> None:
                 self.logger.info("epoch=%d step=%d role=policy stage=train_end rows=%d", epoch, step, len(train_batch))
                 record_dashboard_state(self.areno, stage="train_end", epoch=epoch, step=step, role="policy")
                 self.logger.info("epoch=%d step=%d train_stats=%s", epoch, step, result)
+                if batch_mix_counts:
+                    mix_progress["trained"].update(batch_mix_counts)
+                    mix_progress["target_tokens"].update(batch_mix_target_tokens)
+                    self.logger.info(
+                        "epoch=%d step=%d stage=dataset_mix_progress dataset_mix=%s",
+                        epoch,
+                        step,
+                        _dataset_mix_progress(mix_progress, mix_source_names),
+                    )
                 self._maybe_save(epoch, step)
                 step += 1
                 if self.config.max_steps is not None and step >= self.config.max_steps:
                     self.logger.info("epoch=%d step=%d stage=max_steps_reached", epoch, step)
                     record_dashboard_state(self.areno, stage="max_steps_reached", epoch=epoch, step=step, role="policy")
                     return
+            if mix_source_names:
+                self.logger.info(
+                    "epoch=%d stage=dataset_mix_epoch_end dataset_mix=%s",
+                    epoch,
+                    _dataset_mix_progress(mix_progress, mix_source_names),
+                )
             self.logger.info("epoch=%d stage=epoch_end", epoch)
             record_dashboard_state(self.areno, stage="epoch_end", epoch=epoch, step=step, role="policy")
 
-    def _iter_train_batches(self, tokenizer, processor, *, max_prompt_tokens: int, max_new_tokens: int):
+    def _iter_train_batches(
+        self,
+        tokenizer,
+        processor,
+        *,
+        max_prompt_tokens: int,
+        max_new_tokens: int,
+        mix_progress: dict[str, Counter[str]] | None = None,
+    ):
         # Dataset rows are converted lazily so large HF datasets do not need an
         # up-front tokenized copy. Rows that are empty, all-prompt, or exceed
         # the configured prompt or supervised-response budgets are dropped.
         batch = []
+        batch_mix_counts: Counter[str] = Counter()
+        batch_mix_target_tokens: Counter[str] = Counter()
         skipped = 0
         accepted = 0
         total_rows = len(self.dataset)
         for index in range(total_rows):
             # Normalize each supported row schema into one TrainSequence.
+            record = self.dataset[index]
+            source_name = _dataset_mix_source_name(record)
+            if source_name is not None and mix_progress is not None:
+                mix_progress["scheduled"][source_name] += 1
             seq = _record_to_train_sequence(
-                self.dataset[index],
+                record,
                 tokenizer,
                 processor,
                 max_prompt_tokens=max_prompt_tokens,
@@ -123,12 +180,19 @@ def _iter_train_batches(self, tokenizer, processor, *, max_prompt_tokens: int, m
             )
             if seq is None:
                 skipped += 1
+                if source_name is not None and mix_progress is not None:
+                    mix_progress["filtered"][source_name] += 1
                 continue
             accepted += 1
             batch.append(seq)
+            if source_name is not None:
+                batch_mix_counts[source_name] += 1
+                batch_mix_target_tokens[source_name] += _sft_target_token_count(seq)
             if len(batch) >= self.config.batch_size:
-                yield batch
+                yield batch, batch_mix_counts, batch_mix_target_tokens
                 batch = []
+                batch_mix_counts = Counter()
+                batch_mix_target_tokens = Counter()
         if skipped:
             self.logger.info("stage=sft_dataset_filter skipped_long_or_empty=%d", skipped)
         if accepted == 0:
@@ -138,7 +202,7 @@ def _iter_train_batches(self, tokenizer, processor, *, max_prompt_tokens: int, m
                 "Check dataset quality, --max-prompt-tokens, and --max-new-tokens."
             )
         if batch:
-            yield batch
+            yield batch, batch_mix_counts, batch_mix_target_tokens
 
     def _maybe_save(self, epoch: int, step: int) -> None:
         # Keep the same step-based checkpoint cadence as the RL trainers.
@@ -270,4 +334,53 @@ def _record_to_train_sequence(record: Any, tokenizer, processor=None, *, max_pro
     )
 
 
+def _dataset_mix_source_name(record: Any) -> str | None:
+    if not isinstance(record, Mapping):
+        return None
+    metadata = record.get(DATASET_MIX_METADATA_KEY)
+    if not isinstance(metadata, Mapping):
+        return None
+    source_name = metadata.get("source")
+    return source_name if isinstance(source_name, str) else None
+
+
+def _sft_target_token_count(sequence: areno.api.TrainSequence) -> int:
+    """Count positions that contribute to SFT loss after next-token alignment."""
+
+    prompt_mask = sequence.prompt_mask[1:]
+    if sequence.loss_mask:
+        return sum(
+            not is_prompt and is_enabled
+            for is_prompt, is_enabled in zip(prompt_mask, sequence.loss_mask[1:], strict=True)
+        )
+    return prompt_mask.count(False)
+
+
+def _dataset_mix_progress(progress: dict[str, Counter[str]], source_names: list[str]) -> dict[str, Any]:
+    rows_scheduled = sum(progress["scheduled"].values())
+    rows_filtered = sum(progress["filtered"].values())
+    rows_trained = sum(progress["trained"].values())
+    target_tokens = sum(progress["target_tokens"].values())
+    return {
+        "rows_scheduled": rows_scheduled,
+        "rows_filtered": rows_filtered,
+        "rows_trained": rows_trained,
+        "target_tokens_trained": target_tokens,
+        "sources": [
+            {
+                "name": name,
+                "rows_scheduled": progress["scheduled"][name],
+                "rows_filtered": progress["filtered"][name],
+                "rows_trained": progress["trained"][name],
+                "target_tokens_trained": progress["target_tokens"][name],
+                "observed_sample_proportion": progress["trained"][name] / rows_trained if rows_trained else 0.0,
+                "observed_token_proportion": (
+                    progress["target_tokens"][name] / target_tokens if target_tokens else 0.0
+                ),
+            }
+            for name in sorted(source_names)
+        ],
+    }
+
+
 __all__ = ["SFTTrainer"]
diff --git a/areno/cli/train.py b/areno/cli/train.py
index 1de993a3..1dd1f0a3 100644
--- a/areno/cli/train.py
+++ b/areno/cli/train.py
@@ -17,8 +17,11 @@
 import importlib.util
 import json
 import logging
+import math
+import os
 import shutil
 import textwrap
+from collections.abc import Mapping
 from dataclasses import asdict, fields, is_dataclass
 from pathlib import Path
 from types import SimpleNamespace
@@ -27,6 +30,8 @@
 import click
 
 from areno.api.algorithms import get_algorithm
+from areno.api.data import DATASET_MIX_METADATA_KEY, DatasetMixSource, WeightedMixedDataset
+from areno.api.dataset_mix_artifacts import write_dataset_mix_plan
 from areno.api.defaults import DEFAULT_METRICS_LOG_DIR
 from areno.api.trainer_config import (
     DPOTrainerConfig,
@@ -84,6 +89,11 @@ def flash_attention_unsupported_model_reason(model_config):
             "ckpt",
             "base_model_name_or_path",
             "dataset_path",
+            "dataset_mix_config",
+            "dataset_sources",
+            "dataset_mix_seed",
+            "dataset_mix_exhaustion",
+            "dataset_mix_samples_per_epoch",
             "model_hub",
             "dataset_loader_fn",
             "tune_params",
@@ -233,6 +243,12 @@ def _trainer_config_from_options(**options) -> TrainerConfig:
     args.max_steps = getattr(args, "max_steps", None)
     args.score_micro_bs = getattr(args, "score_micro_bs", 8)
     args.model_hub = getattr(args, "model_hub", "modelscope")
+    dataset_mix_config = getattr(args, "dataset_mix_config", None)
+    args.dataset_mix_config = str(dataset_mix_config) if dataset_mix_config is not None else None
+    args.dataset_sources = tuple(getattr(args, "dataset_sources", ()))
+    args.dataset_mix_seed = getattr(args, "dataset_mix_seed", 42)
+    args.dataset_mix_exhaustion = getattr(args, "dataset_mix_exhaustion", "cycle")
+    args.dataset_mix_samples_per_epoch = getattr(args, "dataset_mix_samples_per_epoch", None)
     args.base_model_name_or_path = getattr(args, "base_model_name_or_path", None)
     args.train_devices = getattr(args, "train_devices", None)
     args.sequence_parallel = getattr(args, "sequence_parallel", None)
@@ -285,13 +301,43 @@ def _trainer_config_from_options(**options) -> TrainerConfig:
     # inputs while RL algorithms still require a reward function or model.
     if args.ckpt is None:
         raise click.UsageError("--ckpt is required")
-    if args.dataset_path is None:
-        raise click.UsageError("--dataset-path is required")
+    dataset_inputs = sum(
+        (
+            args.dataset_path is not None,
+            args.dataset_mix_config is not None,
+            bool(args.dataset_sources),
+        )
+    )
+    if dataset_inputs == 0:
+        raise click.UsageError("one of --dataset-path, --dataset-mix-config, or repeated --dataset-source is required")
+    if dataset_inputs > 1:
+        raise click.UsageError("--dataset-path, --dataset-mix-config, and --dataset-source are mutually exclusive")
+    if not args.dataset_sources and (
+        args.dataset_mix_seed != 42
+        or args.dataset_mix_exhaustion != "cycle"
+        or args.dataset_mix_samples_per_epoch is not None
+    ):
+        raise click.UsageError(
+            "--dataset-mix-seed, --dataset-mix-exhaustion, and --dataset-mix-samples-per-epoch "
+            "apply only to repeated --dataset-source"
+        )
     if args.model_hub not in {"hf", "modelscope"}:
         raise click.UsageError("--model-hub must be one of: hf, modelscope")
     algorithm = _algorithm_for_cli(args.algo)
     if algorithm.name == "sft" and args.dataset_loader_fn is None:
         raise click.UsageError("--dataset-loader-fn is required for --algo sft")
+    if args.dataset_mix_config is not None or args.dataset_sources:
+        if algorithm.name != "sft":
+            raise click.UsageError("dataset mixing currently supports --algo sft only")
+        if args.dataset_mix_config is not None:
+            _preflight_dataset_mix_config(args.dataset_mix_config)
+        else:
+            _preflight_dataset_sources(
+                args.dataset_sources,
+                seed=args.dataset_mix_seed,
+                exhaustion=args.dataset_mix_exhaustion,
+                samples_per_epoch=args.dataset_mix_samples_per_epoch,
+            )
     tune_params = bool(getattr(args, "tune_params", False))
     mem_frac = float(getattr(args, "mem_frac", 0.9))
     tune_max_samples = int(getattr(args, "tune_max_samples", 256))
@@ -409,6 +455,116 @@ def _require_positive_float(value: float, option_name: str) -> None:
         raise click.UsageError(f"{option_name} must be positive")
 
 
+def _preflight_dataset_mix_config(config_path: str) -> None:
+    try:
+        _read_dataset_mix_manifest(config_path)
+    except (OSError, ValueError) as exc:
+        raise click.UsageError(f"invalid --dataset-mix-config: {exc}") from exc
+
+
+def _preflight_dataset_sources(
+    source_specs: tuple[str, ...],
+    *,
+    seed: int = 42,
+    exhaustion: str = "cycle",
+    samples_per_epoch: int | None = None,
+) -> None:
+    try:
+        _dataset_mix_manifest_from_sources(
+            source_specs,
+            seed=seed,
+            exhaustion=exhaustion,
+            samples_per_epoch=samples_per_epoch,
+        )
+    except ValueError as exc:
+        raise click.UsageError(f"invalid --dataset-source: {exc}") from exc
+
+
+def _dataset_mix_manifest_from_sources(
+    source_specs: tuple[str, ...],
+    *,
+    seed: int = 42,
+    exhaustion: str = "cycle",
+    samples_per_epoch: int | None = None,
+) -> dict:
+    if len(source_specs) < 2:
+        raise ValueError("repeat the option at least twice using NAME=PATH:WEIGHT")
+
+    sources = []
+    names: set[str] = set()
+    for index, spec in enumerate(source_specs):
+        source_text, separator, weight_text = spec.rpartition(":")
+        name, name_separator, dataset_path = source_text.partition("=")
+        if (
+            not separator
+            or not name_separator
+            or not name.strip()
+            or not dataset_path.strip()
+            or not weight_text.strip()
+        ):
+            raise ValueError(f"entry {index + 1} must use NAME=PATH:WEIGHT")
+        name = _normalize_dataset_mix_source_name(name, f"entry {index + 1} name")
+        dataset_path = _normalize_dataset_mix_source_path(dataset_path, f"entry {index + 1} path")
+        if name in names:
+            raise ValueError(f"duplicate source name: {name}")
+        names.add(name)
+        try:
+            weight = float(weight_text)
+        except (OverflowError, ValueError):
+            weight = math.nan
+        if not math.isfinite(weight) or weight <= 0:
+            raise ValueError(f"source '{name}' weight must be finite and positive")
+        sources.append({"name": name, "path": dataset_path, "weight": weight})
+
+    _validate_dataset_mix_weight_range([source["weight"] for source in sources], "source weights")
+    if isinstance(seed, bool) or not isinstance(seed, int) or not 0 <= seed < 2**63:
+        raise ValueError("seed must be an integer in [0, 2^63)")
+    if exhaustion not in {"stop", "cycle", "renormalize"}:
+        raise ValueError("exhaustion must be one of: stop, cycle, renormalize")
+    if samples_per_epoch is not None and (
+        isinstance(samples_per_epoch, bool) or not isinstance(samples_per_epoch, int) or samples_per_epoch <= 0
+    ):
+        raise ValueError("samples per epoch must be a positive integer")
+    if samples_per_epoch is not None and exhaustion != "cycle":
+        raise ValueError("samples per epoch requires exhaustion='cycle'")
+    return {
+        "version": 1,
+        "seed": seed,
+        "exhaustion": exhaustion,
+        "shuffle_within_sources": True,
+        "samples_per_epoch": samples_per_epoch,
+        "sources": sources,
+    }
+
+
+def _normalize_dataset_mix_source_name(name: str, label: str) -> str:
+    normalized = name.strip()
+    if not normalized:
+        raise ValueError(f"{label} must be a non-empty string")
+    if not normalized.isprintable():
+        raise ValueError(f"{label} must not contain non-printable characters")
+    return normalized
+
+
+def _normalize_dataset_mix_source_path(dataset_path: str, label: str) -> str:
+    normalized = dataset_path.strip()
+    if not normalized:
+        raise ValueError(f"{label} must be a non-empty string")
+    if not normalized.isprintable():
+        raise ValueError(f"{label} must not contain non-printable characters")
+    return normalized
+
+
+def _validate_dataset_mix_weight_range(weights: list[float], label: str) -> None:
+    max_weight = max(weights)
+    scaled_weights = [weight / max_weight for weight in weights]
+    if any(weight == 0.0 for weight in scaled_weights):
+        raise ValueError(f"{label} have an unsupported numeric range")
+    scaled_total = sum(scaled_weights)
+    if any(weight / scaled_total == 0.0 for weight in scaled_weights):
+        raise ValueError(f"{label} have an unsupported numeric range")
+
+
 def _lora_config_from_options(args):
     rank = getattr(args, "lora_rank", None)
     adapter_path = getattr(args, "lora_adapter_path", None)
@@ -483,7 +639,15 @@ def _format_training_config_summary(
             "Inputs",
             [
                 ("ckpt", config.ckpt),
-                ("dataset_path", config.dataset_path),
+                ("dataset_path", _format_optional(config.dataset_path)),
+                (
+                    "dataset_mix",
+                    (
+                        f"{len(config.dataset_sources)} command-line sources"
+                        if config.dataset_sources
+                        else _format_optional(config.dataset_mix_config)
+                    ),
+                ),
                 ("model_hub", config.model_hub),
                 ("dataset_loader", _format_optional(config.dataset_loader_fn)),
                 (
@@ -841,6 +1005,11 @@ def _trainer_config_from_args(args) -> TrainerConfig:
     args.backend = getattr(args, "backend", None)
     args.score_micro_bs = getattr(args, "score_micro_bs", 8)
     args.model_hub = getattr(args, "model_hub", "modelscope")
+    args.dataset_mix_config = getattr(args, "dataset_mix_config", None)
+    args.dataset_sources = tuple(getattr(args, "dataset_sources", ()))
+    args.dataset_mix_seed = getattr(args, "dataset_mix_seed", 42)
+    args.dataset_mix_exhaustion = getattr(args, "dataset_mix_exhaustion", "cycle")
+    args.dataset_mix_samples_per_epoch = getattr(args, "dataset_mix_samples_per_epoch", None)
     args.base_model_name_or_path = getattr(args, "base_model_name_or_path", None)
     args.train_devices = getattr(args, "train_devices", None)
     args.sequence_parallel = getattr(args, "sequence_parallel", None)
@@ -871,6 +1040,11 @@ def _trainer_config_from_args(args) -> TrainerConfig:
             base_model_name_or_path=args.base_model_name_or_path,
             model_hub=args.model_hub,
             dataset_loader_fn=args.dataset_loader_fn,
+            dataset_mix_config=args.dataset_mix_config,
+            dataset_sources=args.dataset_sources,
+            dataset_mix_seed=args.dataset_mix_seed,
+            dataset_mix_exhaustion=args.dataset_mix_exhaustion,
+            dataset_mix_samples_per_epoch=args.dataset_mix_samples_per_epoch,
             save_path=args.save_path,
             save_interval=args.save_interval,
             epochs=args.epochs,
@@ -931,6 +1105,11 @@ def _trainer_config_from_args(args) -> TrainerConfig:
             base_model_name_or_path=args.base_model_name_or_path,
             model_hub=args.model_hub,
             dataset_loader_fn=args.dataset_loader_fn,
+            dataset_mix_config=args.dataset_mix_config,
+            dataset_sources=args.dataset_sources,
+            dataset_mix_seed=args.dataset_mix_seed,
+            dataset_mix_exhaustion=args.dataset_mix_exhaustion,
+            dataset_mix_samples_per_epoch=args.dataset_mix_samples_per_epoch,
             save_path=args.save_path,
             save_interval=args.save_interval,
             epochs=args.epochs,
@@ -989,6 +1168,11 @@ def _trainer_config_from_args(args) -> TrainerConfig:
             base_model_name_or_path=args.base_model_name_or_path,
             model_hub=args.model_hub,
             dataset_loader_fn=args.dataset_loader_fn,
+            dataset_mix_config=args.dataset_mix_config,
+            dataset_sources=args.dataset_sources,
+            dataset_mix_seed=args.dataset_mix_seed,
+            dataset_mix_exhaustion=args.dataset_mix_exhaustion,
+            dataset_mix_samples_per_epoch=args.dataset_mix_samples_per_epoch,
             reward_fn_path=args.reward_fn_path,
             save_path=args.save_path,
             save_interval=args.save_interval,
@@ -1058,6 +1242,11 @@ def _trainer_config_from_args(args) -> TrainerConfig:
         base_model_name_or_path=args.base_model_name_or_path,
         model_hub=args.model_hub,
         dataset_loader_fn=args.dataset_loader_fn,
+        dataset_mix_config=args.dataset_mix_config,
+        dataset_sources=args.dataset_sources,
+        dataset_mix_seed=args.dataset_mix_seed,
+        dataset_mix_exhaustion=args.dataset_mix_exhaustion,
+        dataset_mix_samples_per_epoch=args.dataset_mix_samples_per_epoch,
         reward_fn_path=args.reward_fn_path,
         save_path=args.save_path,
         save_interval=args.save_interval,
@@ -1146,6 +1335,29 @@ def run(trainer_config: TrainerConfig):
     from areno.api.rewards import load_reward_fn
     from areno.api.trainer_factory import build_trainer
 
+    mixed_dataset = None
+    if trainer_config.dataset_mix_config is not None:
+        mixed_dataset = _load_mixed_dataset_for_training(
+            trainer_config.dataset_mix_config,
+            model_hub=trainer_config.model_hub,
+            dataset_loader_fn=trainer_config.dataset_loader_fn,
+            load_dataset=load_dataset,
+            load_from_disk=load_from_disk,
+        )
+    elif trainer_config.dataset_sources:
+        mixed_dataset = _load_dataset_sources_for_training(
+            trainer_config.dataset_sources,
+            seed=trainer_config.dataset_mix_seed,
+            exhaustion=trainer_config.dataset_mix_exhaustion,
+            samples_per_epoch=trainer_config.dataset_mix_samples_per_epoch,
+            model_hub=trainer_config.model_hub,
+            dataset_loader_fn=trainer_config.dataset_loader_fn,
+            load_dataset=load_dataset,
+            load_from_disk=load_from_disk,
+        )
+    if mixed_dataset is not None:
+        click.echo(_format_dataset_mix_summary(mixed_dataset.summary()))
+
     trainer_config = resolve_model_refs_for_config(trainer_config)
     _write_dashboard_run_config(trainer_config)
     loss_fn = _loss_fn_for_config(trainer_config)
@@ -1160,13 +1372,19 @@ def run(trainer_config: TrainerConfig):
         custom_config=trainer_config.backend_config(),
         score_micro_bs=trainer_config.score_micro_bs,
     )
-    dataset = _load_dataset_for_training(
-        trainer_config.dataset_path,
-        dataset_loader_fn=trainer_config.dataset_loader_fn,
-        model_hub=trainer_config.model_hub,
-        load_dataset=load_dataset,
-        load_from_disk=load_from_disk,
-    )
+    if mixed_dataset is None:
+        if trainer_config.dataset_path is None:
+            raise ValueError("dataset_path is required when dataset mixing is not configured")
+        dataset = _load_dataset_for_training(
+            trainer_config.dataset_path,
+            dataset_loader_fn=trainer_config.dataset_loader_fn,
+            model_hub=trainer_config.model_hub,
+            load_dataset=load_dataset,
+            load_from_disk=load_from_disk,
+        )
+    else:
+        dataset = mixed_dataset
+        _write_dataset_mix_artifact(dataset, trainer_config.metrics_log_dir)
     trainer = build_trainer(trainer_config, instance=api_trainer, dataset=dataset, reward_fn=reward_fn, loss_fn=loss_fn)
     trainer.fit()
 
@@ -1176,7 +1394,6 @@ def _write_dashboard_run_config(config: TrainerConfig) -> None:
 
     if not config.metrics_log_dir:
         return
-    import os
 
     path = Path(config.metrics_log_dir)
     path.mkdir(parents=True, exist_ok=True)
@@ -1217,6 +1434,11 @@ def section(title: str, names: list[str]) -> dict:
                 "algo",
                 "ckpt",
                 "dataset_path",
+                "dataset_mix_config",
+                "dataset_sources",
+                "dataset_mix_seed",
+                "dataset_mix_exhaustion",
+                "dataset_mix_samples_per_epoch",
                 "model_hub",
                 "dataset_loader_fn",
                 "epochs",
@@ -1344,7 +1566,13 @@ def _reward_fn_path_for_config(config: TrainerConfig) -> str | None:
 
 
 def _load_dataset_for_training(
-    dataset_path: str, *, model_hub: str = "modelscope", dataset_loader_fn: str | None, load_dataset, load_from_disk
+    dataset_path: str,
+    *,
+    model_hub: str = "modelscope",
+    dataset_loader_fn: str | None,
+    load_dataset,
+    load_from_disk,
+    _loader_fn=None,
 ):
     def default_loader(path):
         return _load_dataset_from_path(
@@ -1354,8 +1582,10 @@ def default_loader(path):
             load_from_disk=load_from_disk,
         )
 
-    if dataset_loader_fn is not None:
+    loader_fn = _loader_fn
+    if loader_fn is None and dataset_loader_fn is not None:
         loader_fn = _load_dataset_loader_fn(dataset_loader_fn)
+    if loader_fn is not None:
         return loader_fn(
             dataset_path,
             default_loader=default_loader,
@@ -1365,6 +1595,263 @@ def default_loader(path):
     return default_loader(dataset_path)
 
 
+def _read_dataset_mix_manifest(config_path: str | Path) -> dict:
+    path = Path(config_path)
+    try:
+        payload = json.loads(path.read_text(encoding="utf-8"))
+    except json.JSONDecodeError as exc:
+        raise ValueError(f"{path}: invalid JSON at line {exc.lineno}, column {exc.colno}") from exc
+    if not isinstance(payload, dict):
+        raise ValueError(f"{path}: top-level value must be an object")
+    allowed_fields = {
+        "version",
+        "seed",
+        "exhaustion",
+        "shuffle_within_sources",
+        "samples_per_epoch",
+        "max_samples_per_epoch",
+        "sources",
+    }
+    unknown_fields = sorted(set(payload) - allowed_fields)
+    if unknown_fields:
+        raise ValueError(f"{path}: unsupported field(s): {', '.join(unknown_fields)}")
+    version = payload.get("version")
+    if isinstance(version, bool) or not isinstance(version, int) or version != 1:
+        raise ValueError(f"{path}: version must be 1")
+
+    seed = payload.get("seed")
+    if isinstance(seed, bool) or not isinstance(seed, int) or not 0 <= seed < 2**63:
+        raise ValueError(f"{path}: seed must be an integer in [0, 2^63)")
+    exhaustion = payload.get("exhaustion")
+    if exhaustion not in {"stop", "cycle", "renormalize"}:
+        raise ValueError(f"{path}: exhaustion must be one of: stop, cycle, renormalize")
+    shuffle = payload.get("shuffle_within_sources", True)
+    if not isinstance(shuffle, bool):
+        raise ValueError(f"{path}: shuffle_within_sources must be a boolean")
+    if "samples_per_epoch" in payload and "max_samples_per_epoch" in payload:
+        raise ValueError(f"{path}: samples_per_epoch and legacy max_samples_per_epoch are mutually exclusive")
+    samples_per_epoch = payload.get("samples_per_epoch", payload.get("max_samples_per_epoch"))
+    if samples_per_epoch is not None and (
+        isinstance(samples_per_epoch, bool) or not isinstance(samples_per_epoch, int) or samples_per_epoch <= 0
+    ):
+        raise ValueError(f"{path}: samples_per_epoch must be a positive integer")
+    if samples_per_epoch is not None and exhaustion != "cycle":
+        raise ValueError(f"{path}: samples_per_epoch is only supported with exhaustion='cycle'")
+
+    sources = payload.get("sources")
+    if not isinstance(sources, list) or len(sources) < 2:
+        raise ValueError(f"{path}: sources must contain at least two entries")
+    names: set[str] = set()
+    normalized_sources = []
+    for index, source in enumerate(sources):
+        prefix = f"{path}: sources[{index}]"
+        if not isinstance(source, dict):
+            raise ValueError(f"{prefix} must be an object")
+        unknown_source_fields = sorted(set(source) - {"name", "path", "weight"})
+        if unknown_source_fields:
+            raise ValueError(f"{prefix} has unsupported field(s): {', '.join(unknown_source_fields)}")
+        name = source.get("name")
+        dataset_path = source.get("path")
+        weight = source.get("weight")
+        if not isinstance(name, str):
+            raise ValueError(f"{prefix}.name must be a non-empty string")
+        name = _normalize_dataset_mix_source_name(name, f"{prefix}.name")
+        if name in names:
+            raise ValueError(f"{path}: duplicate source name: {name}")
+        names.add(name)
+        if not isinstance(dataset_path, str):
+            raise ValueError(f"{prefix}.path must be a non-empty string")
+        dataset_path = _normalize_dataset_mix_source_path(dataset_path, f"{prefix}.path")
+        try:
+            numeric_weight = float(weight)
+        except (OverflowError, TypeError, ValueError):
+            numeric_weight = math.nan
+        if isinstance(weight, bool) or not math.isfinite(numeric_weight) or numeric_weight <= 0:
+            raise ValueError(f"{prefix}.weight must be finite and positive")
+        normalized_sources.append({"name": name, "path": dataset_path, "weight": numeric_weight})
+    try:
+        _validate_dataset_mix_weight_range(
+            [source["weight"] for source in normalized_sources],
+            "source weights",
+        )
+    except ValueError as exc:
+        raise ValueError(f"{path}: {exc}") from exc
+
+    return {
+        "version": 1,
+        "seed": seed,
+        "exhaustion": exhaustion,
+        "shuffle_within_sources": shuffle,
+        "samples_per_epoch": samples_per_epoch,
+        "sources": normalized_sources,
+    }
+
+
+def _load_mixed_dataset_for_training(
+    config_path: str,
+    *,
+    model_hub: str,
+    dataset_loader_fn: str | None,
+    load_dataset,
+    load_from_disk,
+) -> WeightedMixedDataset:
+    manifest = _read_dataset_mix_manifest(config_path)
+    return _load_mixed_dataset_from_manifest(
+        manifest,
+        source_path_resolver=lambda source_path: _resolve_mix_source_path(Path(config_path), source_path),
+        model_hub=model_hub,
+        dataset_loader_fn=dataset_loader_fn,
+        load_dataset=load_dataset,
+        load_from_disk=load_from_disk,
+    )
+
+
+def _load_dataset_sources_for_training(
+    source_specs: tuple[str, ...],
+    *,
+    seed: int = 42,
+    exhaustion: str = "cycle",
+    samples_per_epoch: int | None = None,
+    model_hub: str,
+    dataset_loader_fn: str | None,
+    load_dataset,
+    load_from_disk,
+) -> WeightedMixedDataset:
+    manifest = _dataset_mix_manifest_from_sources(
+        source_specs,
+        seed=seed,
+        exhaustion=exhaustion,
+        samples_per_epoch=samples_per_epoch,
+    )
+    return _load_mixed_dataset_from_manifest(
+        manifest,
+        source_path_resolver=lambda source_path: source_path,
+        model_hub=model_hub,
+        dataset_loader_fn=dataset_loader_fn,
+        load_dataset=load_dataset,
+        load_from_disk=load_from_disk,
+    )
+
+
+def _load_mixed_dataset_from_manifest(
+    manifest: dict,
+    *,
+    source_path_resolver,
+    model_hub: str,
+    dataset_loader_fn: str | None,
+    load_dataset,
+    load_from_disk,
+) -> WeightedMixedDataset:
+    sources = []
+    try:
+        loader_fn = _load_dataset_loader_fn(dataset_loader_fn) if dataset_loader_fn is not None else None
+    except Exception as exc:
+        raise ValueError(f"stage=dataset_mix_loader input={dataset_loader_fn}: {exc}") from exc
+    for source in manifest["sources"]:
+        source_path = source_path_resolver(source["path"])
+        try:
+            dataset = _load_dataset_for_training(
+                source_path,
+                model_hub=model_hub,
+                dataset_loader_fn=dataset_loader_fn,
+                load_dataset=load_dataset,
+                load_from_disk=load_from_disk,
+                _loader_fn=loader_fn,
+            )
+            _validate_sft_mix_source(source["name"], dataset)
+        except Exception as exc:
+            raise ValueError(
+                f"stage=dataset_mix_validation source={source['name']} input={source['path']}: {exc}"
+            ) from exc
+        sources.append(DatasetMixSource(name=source["name"], dataset=dataset, weight=source["weight"]))
+    samples_per_epoch = manifest["samples_per_epoch"]
+    if manifest["exhaustion"] == "cycle" and samples_per_epoch is None:
+        samples_per_epoch = sum(len(source.dataset) for source in sources)
+    return WeightedMixedDataset(
+        sources,
+        seed=manifest["seed"],
+        exhaustion=manifest["exhaustion"],
+        shuffle_within_sources=manifest["shuffle_within_sources"],
+        samples_per_epoch=samples_per_epoch,
+    )
+
+
+def _resolve_mix_source_path(config_path: Path, source_path: str) -> str:
+    path = Path(source_path)
+    if path.is_absolute():
+        return str(path)
+    candidate = config_path.resolve().parent / path
+    if (
+        candidate.exists()
+        or source_path.startswith(("./", "../"))
+        or path.suffix.lower() in _SUPPORTED_DATASET_SUFFIXES
+    ):
+        return str(candidate)
+    return source_path
+
+
+def _validate_sft_mix_source(source_name: str, dataset) -> None:
+    if not hasattr(dataset, "__len__") or not hasattr(dataset, "__getitem__"):
+        raise ValueError("source must be a map-style dataset with len() and indexed access")
+    if len(dataset) == 0:
+        raise ValueError("source is empty")
+    columns = getattr(dataset, "column_names", None)
+    if columns is not None:
+        _validate_sft_mix_fields(set(columns))
+        if DATASET_MIX_METADATA_KEY in columns:
+            raise ValueError(f"source contains reserved field '{DATASET_MIX_METADATA_KEY}'")
+    rows = [dataset[0]] if columns is not None else (dataset[index] for index in range(len(dataset)))
+    for index, row in enumerate(rows):
+        if not isinstance(row, Mapping):
+            raise ValueError(f"row {index} must be a mapping")
+        try:
+            _validate_sft_mix_fields(set(row))
+        except ValueError as exc:
+            raise ValueError(f"row {index} {exc}") from exc
+        if DATASET_MIX_METADATA_KEY in row:
+            raise ValueError(f"row {index} contains reserved field '{DATASET_MIX_METADATA_KEY}'")
+
+
+def _validate_sft_mix_fields(fields: set[str]) -> None:
+    supported = (
+        {"prompt", "response"}.issubset(fields)
+        or {"tokens", "prompt_mask"}.issubset(fields)
+        or ("response" in fields and bool({"image_base64", "images_base64"} & fields))
+    )
+    if not supported:
+        raise ValueError(
+            "missing required SFT field(s) for a supported row schema: "
+            "prompt+response, tokens+prompt_mask, or image_base64/images_base64+response"
+        )
+
+
+def _format_dataset_mix_summary(summary: dict) -> str:
+    lines = [
+        "AReno dataset mix",
+        f"  policy: {summary['policy']}",
+        f"  weight_unit: {summary['weight_unit']}",
+        f"  seed: {summary['seed']}",
+        f"  shuffle_within_sources: {summary['shuffle_within_sources']}",
+        f"  samples_per_epoch: {summary['samples_per_epoch'] or 'natural exhaustion'}",
+        f"  planned_rows: {summary['planned_rows']}",
+        f"  termination_reason: {summary['termination_reason']}",
+        "  sources:",
+    ]
+    for source in summary["sources"]:
+        lines.append(
+            "    - "
+            f"{source['name']}: rows={source['rows_available']} "
+            f"weight={source['weight_requested']:.6f} selected={source['rows_selected']}"
+        )
+    for warning in summary["warnings"]:
+        lines.append(f"  warning: {warning}")
+    return "\n".join(lines)
+
+
+def _write_dataset_mix_artifact(dataset: WeightedMixedDataset, metrics_log_dir: str | None) -> Path | None:
+    return write_dataset_mix_plan(dataset.summary(), metrics_log_dir)
+
+
 def _load_dataset_loader_fn(spec_text: str):
     loader_path, fn_name = _split_loader_fn_spec(spec_text)
     spec = importlib.util.spec_from_file_location(f"areno_example_dataset_loader_{abs(hash(loader_path))}", loader_path)
@@ -1525,6 +2012,39 @@ def _dataset_builder_for_suffix(suffix: str) -> str:
 @click.option(
     "--dataset-path", default=None, help="Training dataset path, HF save_to_disk directory, or remote dataset ref."
 )
+@click.option(
+    "--dataset-mix-config",
+    type=click.Path(path_type=Path, dir_okay=False),
+    default=None,
+    help="JSON manifest for deterministic weighted SFT dataset mixing.",
+)
+@click.option(
+    "--dataset-source",
+    "dataset_sources",
+    multiple=True,
+    metavar="NAME=PATH:WEIGHT",
+    help=("Weighted SFT dataset source; repeat at least twice. Weights are per sampled row."),
+)
+@click.option(
+    "--dataset-mix-seed",
+    type=click.IntRange(min=0, max=2**63 - 1),
+    default=42,
+    show_default=True,
+    help="Seed for command-line dataset-source sampling and source shuffling.",
+)
+@click.option(
+    "--dataset-mix-exhaustion",
+    type=click.Choice(["cycle", "stop", "renormalize"], case_sensitive=False),
+    default="cycle",
+    show_default=True,
+    help="Behavior when a command-line dataset source is exhausted.",
+)
+@click.option(
+    "--dataset-mix-samples-per-epoch",
+    type=click.IntRange(min=1),
+    default=None,
+    help="Sample budget for a cycle mix; defaults to the sum of source row counts.",
+)
 @click.option(
     "--model-hub",
     type=click.Choice(["hf", "modelscope"], case_sensitive=False),
diff --git a/areno/engine/layers/vocab.py b/areno/engine/layers/vocab.py
index df129ba1..4c234540 100644
--- a/areno/engine/layers/vocab.py
+++ b/areno/engine/layers/vocab.py
@@ -42,6 +42,9 @@ def __init__(self, vocab_size: int, hidden_size: int, *, dtype: torch.dtype | No
         self.vocab_start, self.vocab_end = _shard_range(vocab_size, ctx.rank, ctx.world_size)
         self.weight = nn.Parameter(torch.empty(self.vocab_end - self.vocab_start, hidden_size, dtype=dtype))
         mark_tensor_parallel_parameter(self.weight, True, sequence_parallel=True)
+        # Optimizer-role metadata is consumed by AdamW8bit without relying on
+        # module/parameter name substrings. It does not change model outputs.
+        self.weight._areno_optimizer_role = "token_embedding"
         nn.init.normal_(self.weight, mean=0.0, std=0.02)
 
     def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
diff --git a/areno/engine/optim/__init__.py b/areno/engine/optim/__init__.py
index cd27ed45..f2220123 100644
--- a/areno/engine/optim/__init__.py
+++ b/areno/engine/optim/__init__.py
@@ -6,7 +6,7 @@
 """
 
 from areno.engine.optim.adamw_4bit import AdamW4bit
-from areno.engine.optim.adamw_8bit import AdamW8bit
+from areno.engine.optim.adamw_8bit import AdamW8bit, set_optimizer_state_precision
 from areno.engine.optim.adamw_fp32_master import AdamWFP32Master
 
-__all__ = ["AdamW4bit", "AdamW8bit", "AdamWFP32Master"]
+__all__ = ["AdamW4bit", "AdamW8bit", "AdamWFP32Master", "set_optimizer_state_precision"]
diff --git a/areno/engine/optim/adamw_4bit.py b/areno/engine/optim/adamw_4bit.py
index 939da6ad..943bca45 100644
--- a/areno/engine/optim/adamw_4bit.py
+++ b/areno/engine/optim/adamw_4bit.py
@@ -49,6 +49,13 @@ class AdamW4bit(AdamW8bit):
     state directly with a fused block-wise kernel.
     """
 
+    _embedding_fp32_state = False
+    state_quantizer = "signed-de4/zero-excluding-linear4"
+
+    def _precision_for_parameter(self, parameter: torch.nn.Parameter) -> str:
+        del parameter
+        return "8bit"
+
     def __init__(
         self,
         params: Iterable[torch.nn.Parameter],
@@ -82,6 +89,14 @@ def state_dict(self) -> dict:
 
         payload = super().state_dict()
         payload.pop("adam_8bit", None)
+        payload.pop("quantizer", None)
+        payload.pop("precision_policy", None)
+        payload.pop("state_memory", None)
+        for state in payload["state"]:
+            state.pop("precision", None)
+            state.pop("quantizer", None)
+            state.pop("exp_avg", None)
+            state.pop("exp_avg_sq", None)
         payload["adam_4bit"] = True
         payload["state_format_version"] = _STATE_FORMAT_VERSION
         payload["quant_block_size"] = self.quant_block_size
diff --git a/areno/engine/optim/adamw_8bit.py b/areno/engine/optim/adamw_8bit.py
index 632b2136..499efc29 100644
--- a/areno/engine/optim/adamw_8bit.py
+++ b/areno/engine/optim/adamw_8bit.py
@@ -2,8 +2,9 @@
 
 from __future__ import annotations
 
-from collections.abc import Iterable
+from collections.abc import Iterable, Mapping
 from dataclasses import dataclass
+from typing import Any
 
 import torch
 import torch.distributed as dist
@@ -17,9 +18,18 @@
     _param_grad,
     _ParamRef,
 )
+from areno.engine.optim.dynamic_quant import (
+    SIGNED_DYNAMIC_MAP,
+    SIGNED_DYNAMIC_ZERO,
+    UNSIGNED_DYNAMIC_MAP,
+    UNSIGNED_DYNAMIC_ZERO,
+)
 
 _DEFAULT_QUANT_BLOCK_SIZE = 128
 _MAX_FUSED_QUANT_BLOCK_SIZE = 4096
+_DYNAMIC_QUANTIZER = "dynamic-tree-v1"
+_VALID_STATE_PRECISIONS = frozenset({"8bit", "fp32"})
+_CODEBOOK_CACHE: dict[tuple[torch.device, bool], torch.Tensor] = {}
 
 
 @dataclass(slots=True)
@@ -31,6 +41,10 @@ class _Adam8bitBucketState:
     exp_avg_scale: torch.Tensor | None = None
     exp_avg_sq_q: torch.Tensor | None = None
     exp_avg_sq_scale: torch.Tensor | None = None
+    exp_avg: torch.Tensor | None = None
+    exp_avg_sq: torch.Tensor | None = None
+    precision: str = "8bit"
+    quantizer: str = _DYNAMIC_QUANTIZER
     offload_file: str | None = None
     offload_index: int | None = None
     offload_group: _MmapGroup | None = None
@@ -42,12 +56,16 @@ class AdamW8bit(AdamWFP32Master):
 
     The model parameters remain BF16 on every DP rank. Adam moments are stored
     for only this rank's DP shard and re-quantized after every bucket update.
-    This trades optimizer precision for much lower persistent optimizer memory.
+    Explicit token-embedding parameters keep FP32 moments for stability, while
+    other parameters use the paper-compatible dynamic 8-bit codebooks.
     """
 
+    _embedding_fp32_state = True
+    state_quantizer = _DYNAMIC_QUANTIZER
+
     def __init__(
         self,
-        params: Iterable[torch.nn.Parameter],
+        params: Iterable[torch.nn.Parameter] | Iterable[Mapping[str, Any]],
         *,
         lr: float,
         betas: tuple[float, float],
@@ -62,8 +80,12 @@ def __init__(
             raise ValueError(
                 f"quant_block_size must be between 1 and {_MAX_FUSED_QUANT_BLOCK_SIZE}, got {quant_block_size}"
             )
+        normalized_params, precision_by_id, role_by_id = _normalize_parameter_policies(params)
+        self._parameter_state_precision = precision_by_id
+        self._parameter_roles = role_by_id
+        self._bucket_numel = max(bucket_numel, 1)
         super().__init__(
-            params,
+            normalized_params,
             lr=lr,
             betas=betas,
             weight_decay=weight_decay,
@@ -73,7 +95,40 @@ def __init__(
             dp_group=dp_group,
         )
         self.quant_block_size = quant_block_size
-        self._states = [_Adam8bitBucketState() for _ in self.buckets]
+        self._states = [
+            _Adam8bitBucketState(
+                precision=self._precision_for_parameter(bucket.refs[0].model_param),
+                quantizer=self.state_quantizer,
+            )
+            for bucket in self.buckets
+        ]
+
+    @torch.no_grad()
+    def _build_buckets(self, params: list[torch.nn.Parameter], bucket_numel: int) -> list[_MasterBucket]:
+        """Keep FP32-exempt and quantized parameters in separate buckets."""
+
+        buckets: list[_MasterBucket] = []
+        pending: list[torch.nn.Parameter] = []
+        pending_precision: str | None = None
+        for parameter in params:
+            precision = self._precision_for_parameter(parameter)
+            if pending and precision != pending_precision:
+                buckets.extend(AdamWFP32Master._build_buckets(self, pending, bucket_numel))
+                pending = []
+            pending.append(parameter)
+            pending_precision = precision
+        if pending:
+            buckets.extend(AdamWFP32Master._build_buckets(self, pending, bucket_numel))
+        return buckets
+
+    def _precision_for_parameter(self, parameter: torch.nn.Parameter) -> str:
+        explicit = self._parameter_state_precision.get(id(parameter))
+        if explicit is not None:
+            return explicit
+        role = self._parameter_roles.get(id(parameter), getattr(parameter, "_areno_optimizer_role", None))
+        if self._embedding_fp32_state and role == "token_embedding":
+            return "fp32"
+        return "8bit"
 
     @torch.no_grad()
     def step(self, closure=None):
@@ -115,6 +170,8 @@ def clear_state(self) -> None:
             state.exp_avg_scale = None
             state.exp_avg_sq_q = None
             state.exp_avg_sq_scale = None
+            state.exp_avg = None
+            state.exp_avg_sq = None
             state.offload_file = None
             state.offload_index = None
             state.offload_group = None
@@ -168,6 +225,10 @@ def onload_state(self, device: torch.device) -> None:
                 state.exp_avg_sq_q = state.exp_avg_sq_q.to(device=device)
             if state.exp_avg_sq_scale is not None and state.exp_avg_sq_scale.device != device:
                 state.exp_avg_sq_scale = state.exp_avg_sq_scale.to(device=device)
+            if state.exp_avg is not None and state.exp_avg.device != device:
+                state.exp_avg = state.exp_avg.to(device=device)
+            if state.exp_avg_sq is not None and state.exp_avg_sq.device != device:
+                state.exp_avg_sq = state.exp_avg_sq.to(device=device)
             state.offload_file = None
             state.offload_index = None
             state.offload_group = None
@@ -190,14 +251,29 @@ def state_dict(self) -> dict:
             "dp_rank": self.dp_rank,
             "dp_size": self.dp_size,
             "adam_8bit": True,
+            "quantizer": self.state_quantizer,
             "quant_block_size": self.quant_block_size,
+            "precision_policy": [
+                {
+                    "parameter_index": index,
+                    "role": self._parameter_roles.get(id(parameter)),
+                    "precision": self._precision_for_parameter(parameter),
+                    "numel": parameter.numel(),
+                }
+                for index, parameter in enumerate(self.model_params)
+            ],
+            "state_memory": self.state_memory_metrics(),
             "state": [
                 {
                     "step": state.step,
+                    "precision": state.precision,
+                    "quantizer": state.quantizer,
                     "exp_avg_q": payload["exp_avg_q"],
                     "exp_avg_scale": payload["exp_avg_scale"],
                     "exp_avg_sq_q": payload["exp_avg_sq_q"],
                     "exp_avg_sq_scale": payload["exp_avg_sq_scale"],
+                    "exp_avg": payload["exp_avg"],
+                    "exp_avg_sq": payload["exp_avg_sq"],
                 }
                 for state, payload in zip(self._states, payloads, strict=True)
             ],
@@ -211,12 +287,31 @@ def load_state_dict(self, state_dict: dict) -> None:
         self._active_offload_mode = "none"
         self._disk_offload_root = None
         self._active_offload_batch_size = 1
+        saved_states = state_dict.get("state", [])
+        saved_quantizer = str(state_dict.get("quantizer", ""))
+        if saved_quantizer != _DYNAMIC_QUANTIZER:
+            raise ValueError(f"unsupported AdamW8bit quantizer: {saved_quantizer}")
+        if state_dict.get("precision_policy") is None:
+            raise ValueError("AdamW8bit checkpoint is missing its precision policy")
+        self._restore_precision_policy(state_dict["precision_policy"])
+        if len(saved_states) != len(self.buckets):
+            raise ValueError(
+                "AdamW8bit checkpoint bucket count does not match the current optimizer layout: "
+                f"checkpoint={len(saved_states)}, optimizer={len(self.buckets)}"
+            )
         for state in self._states:
+            state.step = 0
+            state.exp_avg_q = None
+            state.exp_avg_scale = None
+            state.exp_avg_sq_q = None
+            state.exp_avg_sq_scale = None
+            state.exp_avg = None
+            state.exp_avg_sq = None
+            state.quantizer = self.state_quantizer
             state.offload_file = None
             state.offload_index = None
             state.offload_group = None
             state.offload_ready_events = ()
-        saved_states = state_dict.get("state", [])
         if "quant_block_size" in state_dict:
             saved_block_size = int(state_dict["quant_block_size"])
             if saved_block_size < 1 or saved_block_size > _MAX_FUSED_QUANT_BLOCK_SIZE:
@@ -227,35 +322,103 @@ def load_state_dict(self, state_dict: dict) -> None:
                 continue
             device = bucket.refs[0].model_param.device
             state.step = int(saved.get("step", 0))
+            saved_precision = str(saved.get("precision", "8bit"))
+            bucket_quantizer = str(saved.get("quantizer", saved_quantizer))
+            if bucket_quantizer != _DYNAMIC_QUANTIZER:
+                raise ValueError(f"unsupported AdamW8bit bucket quantizer: {bucket_quantizer}")
+            if saved_precision != state.precision:
+                raise ValueError(
+                    "AdamW8bit state precision policy mismatch: "
+                    f"checkpoint={saved_precision}, optimizer={state.precision}"
+                )
+            state.quantizer = bucket_quantizer
+            if state.precision == "fp32":
+                state.exp_avg = _load_optional_state_tensor(saved.get("exp_avg"), bucket, device)
+                state.exp_avg_sq = _load_optional_state_tensor(saved.get("exp_avg_sq"), bucket, device)
+                state.exp_avg_q = None
+                state.exp_avg_scale = None
+                state.exp_avg_sq_q = None
+                state.exp_avg_sq_scale = None
+                continue
             exp_avg_q = saved.get("exp_avg_q")
             exp_avg_scale = saved.get("exp_avg_scale")
             exp_avg_sq_q = saved.get("exp_avg_sq_q")
             exp_avg_sq_scale = saved.get("exp_avg_sq_scale")
-            state.exp_avg_q = (
-                None if exp_avg_q is None else exp_avg_q.detach().to(device=device, dtype=torch.uint8).view(-1).clone()
-            )
+            state.exp_avg_q = _load_optional_quantized_tensor(exp_avg_q, bucket, device)
             state.exp_avg_scale = None if exp_avg_scale is None else self._restore_scales(exp_avg_scale, bucket, device)
-            state.exp_avg_sq_q = (
-                None
-                if exp_avg_sq_q is None
-                else exp_avg_sq_q.detach().to(device=device, dtype=torch.uint8).view(-1).clone()
-            )
+            state.exp_avg_sq_q = _load_optional_quantized_tensor(exp_avg_sq_q, bucket, device)
             state.exp_avg_sq_scale = (
                 None if exp_avg_sq_scale is None else self._restore_scales(exp_avg_sq_scale, bucket, device)
             )
 
+    def _restore_precision_policy(self, saved_policy: Any) -> None:
+        """Rebuild buckets from the checkpoint's identity-ordered policy."""
+
+        if not isinstance(saved_policy, list) or len(saved_policy) != len(self.model_params):
+            raise ValueError(
+                "AdamW8bit checkpoint precision policy does not match the current parameter count: "
+                f"checkpoint={len(saved_policy) if isinstance(saved_policy, list) else 'invalid'}, "
+                f"optimizer={len(self.model_params)}"
+            )
+        restored_precision: dict[int, str] = {}
+        restored_roles: dict[int, str] = {}
+        for expected_index, (entry, parameter) in enumerate(zip(saved_policy, self.model_params, strict=True)):
+            if not isinstance(entry, Mapping) or int(entry.get("parameter_index", -1)) != expected_index:
+                raise ValueError(
+                    f"AdamW8bit checkpoint has an invalid precision policy entry at index {expected_index}"
+                )
+            saved_numel = int(entry.get("numel", -1))
+            if saved_numel != parameter.numel():
+                raise ValueError(
+                    "AdamW8bit checkpoint precision policy parameter size mismatch: "
+                    f"index={expected_index}, checkpoint={saved_numel}, optimizer={parameter.numel()}"
+                )
+            restored_precision[id(parameter)] = _normalize_state_precision(str(entry.get("precision", "8bit")))
+            role = entry.get("role")
+            if role is not None:
+                restored_roles[id(parameter)] = str(role)
+        self._parameter_state_precision = restored_precision
+        self._parameter_roles = restored_roles
+        self.buckets = self._build_buckets(self.model_params, self._bucket_numel)
+        self._states = [
+            _Adam8bitBucketState(
+                precision=self._precision_for_parameter(bucket.refs[0].model_param),
+                quantizer=self.state_quantizer,
+            )
+            for bucket in self.buckets
+        ]
+
+    def state_memory_metrics(self) -> dict[str, int]:
+        """Report logical persistent moment storage for initialized buckets."""
+
+        quantized_state_bytes = 0
+        fp32_exempt_bytes = 0
+        block_metadata_bytes = 0
+        for bucket, state in zip(self.buckets, self._states, strict=True):
+            if state.step == 0:
+                continue
+            if state.precision == "fp32":
+                fp32_exempt_bytes += 2 * bucket.shard_numel * 4
+            else:
+                quantized_state_bytes += 2 * bucket.shard_numel
+                block_metadata_bytes += 2 * self._bucket_scale_count(bucket) * 4
+        return {
+            "quantized_state_bytes": quantized_state_bytes,
+            "fp32_exempt_bytes": fp32_exempt_bytes,
+            "block_metadata_bytes": block_metadata_bytes,
+            "total_bytes": quantized_state_bytes + fp32_exempt_bytes + block_metadata_bytes,
+        }
+
     def _restore_scales(
         self,
         saved: torch.Tensor,
         bucket: _MasterBucket,
         device: torch.device,
     ) -> torch.Tensor:
-        """Restore block scales, expanding legacy bucket-level scalar scales."""
+        """Restore block scales for one quantized bucket."""
 
         expected = self._bucket_scale_count(bucket)
         scales = saved.detach().to(device=device, dtype=torch.float32).view(-1)
-        if scales.numel() == 1 and expected != 1:
-            return scales.expand(expected).clone()
         if scales.numel() != expected:
             raise ValueError(f"AdamW8bit checkpoint has {scales.numel()} scales for a bucket requiring {expected}")
         return scales.clone()
@@ -275,12 +438,24 @@ def _ensure_bucket_state(self, bucket: _MasterBucket, state: _Adam8bitBucketStat
             state.exp_avg_sq_q = state.exp_avg_sq_q.to(device=device)
         if state.exp_avg_sq_scale is not None and state.exp_avg_sq_scale.device != device:
             state.exp_avg_sq_scale = state.exp_avg_sq_scale.to(device=device)
+        if state.exp_avg is not None and state.exp_avg.device != device:
+            state.exp_avg = state.exp_avg.to(device=device)
+        if state.exp_avg_sq is not None and state.exp_avg_sq.device != device:
+            state.exp_avg_sq = state.exp_avg_sq.to(device=device)
+        if state.precision == "fp32":
+            if state.exp_avg is None:
+                state.exp_avg = torch.zeros(bucket.shard_numel, device=device, dtype=torch.float32)
+            if state.exp_avg_sq is None:
+                state.exp_avg_sq = torch.zeros(bucket.shard_numel, device=device, dtype=torch.float32)
+            return
         if state.exp_avg_q is None:
-            state.exp_avg_q = torch.full((bucket.shard_numel,), 128, device=device, dtype=torch.uint8)
-            state.exp_avg_scale = torch.ones(self._bucket_scale_count(bucket), device=device, dtype=torch.float32)
+            state.exp_avg_q = torch.full((bucket.shard_numel,), SIGNED_DYNAMIC_ZERO, device=device, dtype=torch.uint8)
+            state.exp_avg_scale = torch.zeros(self._bucket_scale_count(bucket), device=device, dtype=torch.float32)
         if state.exp_avg_sq_q is None:
-            state.exp_avg_sq_q = torch.zeros(bucket.shard_numel, device=device, dtype=torch.uint8)
-            state.exp_avg_sq_scale = torch.ones(self._bucket_scale_count(bucket), device=device, dtype=torch.float32)
+            state.exp_avg_sq_q = torch.full(
+                (bucket.shard_numel,), UNSIGNED_DYNAMIC_ZERO, device=device, dtype=torch.uint8
+            )
+            state.exp_avg_sq_scale = torch.zeros(self._bucket_scale_count(bucket), device=device, dtype=torch.float32)
 
     def _bucket_scale_count(self, bucket: _MasterBucket) -> int:
         """Return the number of independently scaled blocks in one DP shard."""
@@ -309,10 +484,9 @@ def _load_state_offload(self, state: _Adam8bitBucketState, device: torch.device)
             state.offload_index,
             state.offload_group.tensors[state.offload_index],
         )
-        state.exp_avg_q = _host_tensor_to(saved["exp_avg_q"], device, prefetched=prefetched)
-        state.exp_avg_scale = _host_tensor_to(saved["exp_avg_scale"], device, prefetched=prefetched)
-        state.exp_avg_sq_q = _host_tensor_to(saved["exp_avg_sq_q"], device, prefetched=prefetched)
-        state.exp_avg_sq_scale = _host_tensor_to(saved["exp_avg_sq_scale"], device, prefetched=prefetched)
+        for name in ("exp_avg_q", "exp_avg_scale", "exp_avg_sq_q", "exp_avg_sq_scale", "exp_avg", "exp_avg_sq"):
+            value = saved.get(name)
+            setattr(state, name, None if value is None else _host_tensor_to(value, device, prefetched=prefetched))
         if prefetched and device.type == "cuda":
             self._retain_disk_prefetch(state.offload_index, saved, device)
 
@@ -325,22 +499,33 @@ def _disk_mmap_group_for_index(self, index: int) -> _MmapGroup | None:
     def _state_mmap_specs(self, indices: list[int]) -> dict[int, dict[str, tuple[torch.dtype, tuple[int, ...]]]]:
         """Return the fixed raw-mmap layout for quantized Adam state."""
 
-        return {
-            index: {
-                "exp_avg_q": (torch.uint8, (self.buckets[index].shard_numel,)),
-                "exp_avg_scale": (torch.float32, (self._bucket_scale_count(self.buckets[index]),)),
-                "exp_avg_sq_q": (torch.uint8, (self.buckets[index].shard_numel,)),
-                "exp_avg_sq_scale": (torch.float32, (self._bucket_scale_count(self.buckets[index]),)),
-            }
-            for index in indices
-        }
+        specs: dict[int, dict[str, tuple[torch.dtype, tuple[int, ...]]]] = {}
+        for index in indices:
+            bucket = self.buckets[index]
+            if self._states[index].precision == "fp32":
+                specs[index] = {
+                    "exp_avg": (torch.float32, (bucket.shard_numel,)),
+                    "exp_avg_sq": (torch.float32, (bucket.shard_numel,)),
+                }
+            else:
+                specs[index] = {
+                    "exp_avg_q": (torch.uint8, (bucket.shard_numel,)),
+                    "exp_avg_scale": (torch.float32, (self._bucket_scale_count(bucket),)),
+                    "exp_avg_sq_q": (torch.uint8, (bucket.shard_numel,)),
+                    "exp_avg_sq_scale": (torch.float32, (self._bucket_scale_count(bucket),)),
+                }
+        return specs
 
     def _offload_8bit_group_to_disk(self, indices: list[int]) -> None:
         """Persist a bounded group of quantized states in one serialization call."""
 
         if self._disk_offload_root is None:
             raise RuntimeError("disk optimizer offload is active without a usable directory")
-        present_indices = [index for index in indices if self._states[index].exp_avg_q is not None]
+        present_indices = [
+            index
+            for index in indices
+            if self._states[index].exp_avg_q is not None or self._states[index].exp_avg is not None
+        ]
         if not present_indices:
             return
         group = self._get_or_create_mmap_group(indices, self._state_mmap_specs(indices))
@@ -348,15 +533,17 @@ def _offload_8bit_group_to_disk(self, indices: list[int]) -> None:
         ready_events: list[torch.cuda.Event] = []
         for index in present_indices:
             state = self._states[index]
-            assert state.exp_avg_q is not None
-            assert state.exp_avg_scale is not None
-            assert state.exp_avg_sq_q is not None
-            assert state.exp_avg_sq_scale is not None
             payloads[index] = {
-                "exp_avg_q": state.exp_avg_q,
-                "exp_avg_scale": state.exp_avg_scale,
-                "exp_avg_sq_q": state.exp_avg_sq_q,
-                "exp_avg_sq_scale": state.exp_avg_sq_scale,
+                name: value
+                for name in (
+                    "exp_avg_q",
+                    "exp_avg_scale",
+                    "exp_avg_sq_q",
+                    "exp_avg_sq_scale",
+                    "exp_avg",
+                    "exp_avg_sq",
+                )
+                if (value := getattr(state, name)) is not None
             }
             ready_events.extend(state.offload_ready_events)
         self._submit_disk_group_write(indices, group, payloads, tuple(ready_events))
@@ -369,6 +556,8 @@ def _offload_8bit_group_to_disk(self, indices: list[int]) -> None:
             state.exp_avg_scale = None
             state.exp_avg_sq_q = None
             state.exp_avg_sq_scale = None
+            state.exp_avg = None
+            state.exp_avg_sq = None
             state.offload_ready_events = ()
 
     def _stage_8bit_state_on_cpu(self, state: _Adam8bitBucketState) -> None:
@@ -381,6 +570,8 @@ def _stage_8bit_state_on_cpu(self, state: _Adam8bitBucketState) -> None:
                 "exp_avg_scale": state.exp_avg_scale,
                 "exp_avg_sq_q": state.exp_avg_sq_q,
                 "exp_avg_sq_scale": state.exp_avg_sq_scale,
+                "exp_avg": state.exp_avg,
+                "exp_avg_sq": state.exp_avg_sq,
             }.items()
             if tensor is not None
         }
@@ -389,6 +580,8 @@ def _stage_8bit_state_on_cpu(self, state: _Adam8bitBucketState) -> None:
         state.exp_avg_scale = staged.get("exp_avg_scale")
         state.exp_avg_sq_q = staged.get("exp_avg_sq_q")
         state.exp_avg_sq_scale = staged.get("exp_avg_sq_scale")
+        state.exp_avg = staged.get("exp_avg")
+        state.exp_avg_sq = staged.get("exp_avg_sq")
 
     def _state_cpu_payload(
         self,
@@ -402,22 +595,30 @@ def _state_cpu_payload(
             assert state.offload_group is not None
             self._wait_disk_group_write(state.offload_group)
             saved = state.offload_group.tensors[index]
-            return {name: tensor.clone() for name, tensor in saved.items()}
+            return {
+                name: None if saved.get(name) is None else saved[name].clone()
+                for name in (
+                    "exp_avg_q",
+                    "exp_avg_scale",
+                    "exp_avg_sq_q",
+                    "exp_avg_sq_scale",
+                    "exp_avg",
+                    "exp_avg_sq",
+                )
+            }
         return {
             "exp_avg_q": _cpu_clone(state.exp_avg_q),
             "exp_avg_scale": _cpu_clone(state.exp_avg_scale),
             "exp_avg_sq_q": _cpu_clone(state.exp_avg_sq_q),
             "exp_avg_sq_scale": _cpu_clone(state.exp_avg_sq_scale),
+            "exp_avg": _cpu_clone(state.exp_avg),
+            "exp_avg_sq": _cpu_clone(state.exp_avg_sq),
         }
 
     @torch.no_grad()
     def _step_bucket_8bit(self, bucket: _MasterBucket, state: _Adam8bitBucketState) -> None:
         """Update one bucket without materializing full FP32 moment tensors."""
 
-        assert state.exp_avg_q is not None
-        assert state.exp_avg_scale is not None
-        assert state.exp_avg_sq_q is not None
-        assert state.exp_avg_sq_scale is not None
         beta1, beta2 = self.betas
         state.step += 1
         bias_correction1 = 1.0 - beta1**state.step
@@ -430,19 +631,32 @@ def _step_bucket_8bit(self, bucket: _MasterBucket, state: _Adam8bitBucketState)
                 continue
             effective_lr = float(getattr(ref.model_param, "_areno_lr", self.lr))
             step_size = effective_lr / bias_correction1
-            self._step_param_ref_8bit(
-                bucket,
-                ref,
-                grad,
-                state,
-                scale_offset,
-                block_count,
-                beta1,
-                beta2,
-                effective_lr,
-                step_size,
-                bias_correction2_sqrt,
-            )
+            if state.precision == "fp32":
+                self._step_param_ref_fp32(
+                    bucket,
+                    ref,
+                    grad,
+                    state,
+                    beta1,
+                    beta2,
+                    effective_lr,
+                    step_size,
+                    bias_correction2_sqrt,
+                )
+            else:
+                self._step_param_ref_8bit(
+                    bucket,
+                    ref,
+                    grad,
+                    state,
+                    scale_offset,
+                    block_count,
+                    beta1,
+                    beta2,
+                    effective_lr,
+                    step_size,
+                    bias_correction2_sqrt,
+                )
             if ref.param_start + ref.numel == ref.model_param.numel():
                 ref.model_param.grad = None
                 if isinstance(getattr(ref.model_param, "main_grad", None), torch.Tensor):
@@ -453,6 +667,58 @@ def _step_bucket_8bit(self, bucket: _MasterBucket, state: _Adam8bitBucketState)
         self._all_gather_bucket(bucket)
         bucket.grad_shard = None
         bucket.grad_param_ids = frozenset()
+        if state.precision == "8bit":
+            state.quantizer = self.state_quantizer
+
+    @torch.no_grad()
+    def _step_param_ref_fp32(
+        self,
+        bucket: _MasterBucket,
+        ref: _ParamRef,
+        grad: torch.Tensor,
+        state: _Adam8bitBucketState,
+        beta1: float,
+        beta2: float,
+        effective_lr: float,
+        step_size: float,
+        bias_correction2_sqrt: float,
+    ) -> None:
+        """Update one FP32-exempt parameter shard without a master-weight copy."""
+
+        if ref.shard_numel == 0:
+            return
+        assert state.exp_avg is not None
+        assert state.exp_avg_sq is not None
+        grad_shard = grad if bucket.grad_shard is not None else grad.narrow(0, ref.shard_start, ref.shard_numel)
+        model_shard = ref.model_param.detach().reshape(-1).narrow(0, ref.param_start + ref.shard_start, ref.shard_numel)
+        exp_avg = state.exp_avg.narrow(0, ref.shard_bucket_start, ref.shard_numel)
+        exp_avg_sq = state.exp_avg_sq.narrow(0, ref.shard_bucket_start, ref.shard_numel)
+        if model_shard.is_cuda:
+            from areno.accel.optimizer import areno_adamw_fp32_state_step
+
+            areno_adamw_fp32_state_step(
+                model_shard,
+                grad_shard.contiguous(),
+                exp_avg,
+                exp_avg_sq,
+                beta1=beta1,
+                beta2=beta2,
+                effective_lr=effective_lr,
+                weight_decay=self.weight_decay,
+                eps=self.eps,
+                step_size=step_size,
+                bias_correction2_sqrt=bias_correction2_sqrt,
+            )
+            return
+        weight = model_shard.float()
+        gradient = grad_shard.float()
+        if self.weight_decay != 0.0:
+            weight.mul_(1.0 - effective_lr * self.weight_decay)
+        exp_avg.mul_(beta1).add_(gradient, alpha=1.0 - beta1)
+        exp_avg_sq.mul_(beta2).addcmul_(gradient, gradient, value=1.0 - beta2)
+        denom = exp_avg_sq.sqrt().div_(bias_correction2_sqrt).add_(self.eps)
+        weight.addcdiv_(exp_avg, denom, value=-step_size)
+        model_shard.copy_(weight)
 
     @torch.no_grad()
     def _step_param_ref_8bit(
@@ -491,6 +757,9 @@ def _step_param_ref_8bit(
         if model_shard.is_cuda:
             from areno.accel.optimizer import areno_adamw_8bit_step
 
+            signed_codebook = _dynamic_codebook(model_shard.device, signed=True)
+            unsigned_codebook = _dynamic_codebook(model_shard.device, signed=False)
+
             areno_adamw_8bit_step(
                 model_shard,
                 grad_shard.contiguous(),
@@ -498,6 +767,8 @@ def _step_param_ref_8bit(
                 moment_scales,
                 variance_q,
                 variance_scales,
+                signed_codebook,
+                unsigned_codebook,
                 block_size=self.quant_block_size,
                 beta1=beta1,
                 beta2=beta2,
@@ -512,57 +783,173 @@ def _step_param_ref_8bit(
         for block_index in range(block_count):
             start = block_index * self.quant_block_size
             numel = min(self.quant_block_size, ref.shard_numel - start)
-            weight = model_shard.narrow(0, start, numel).to(dtype=torch.float32)
+            # ``Tensor.to(float32)`` aliases an already-FP32 model shard. Keep
+            # this speculative update private until the whole block passes the
+            # finite check, matching the fused CUDA kernel's two-pass commit.
+            weight = model_shard.narrow(0, start, numel).to(dtype=torch.float32).clone()
             block_grad = grad_shard.narrow(0, start, numel).to(dtype=torch.float32)
             block_moment_q = moment_q.narrow(0, start, numel)
             block_variance_q = variance_q.narrow(0, start, numel)
-            moment = _dequantize_symmetric(block_moment_q, moment_scales[block_index])
-            variance = _dequantize_positive(block_variance_q, variance_scales[block_index])
+            moment = _dequantize_dynamic(block_moment_q, moment_scales[block_index], signed=True)
+            variance = _dequantize_dynamic(block_variance_q, variance_scales[block_index], signed=False)
             if self.weight_decay != 0.0:
                 weight.mul_(1.0 - effective_lr * self.weight_decay)
             moment.mul_(beta1).add_(block_grad, alpha=1.0 - beta1)
             variance.mul_(beta2).addcmul_(block_grad, block_grad, value=1.0 - beta2)
             denom = variance.sqrt().div_(bias_correction2_sqrt).add_(self.eps)
             weight.addcdiv_(moment, denom, value=-step_size)
+            if not bool(
+                torch.isfinite(block_grad).all()
+                & torch.isfinite(moment).all()
+                & torch.isfinite(variance).all()
+                & torch.isfinite(weight).all()
+            ):
+                continue
             model_shard.narrow(0, start, numel).copy_(weight)
-            quantized_moment, moment_scale = _quantize_symmetric(moment)
-            quantized_variance, variance_scale = _quantize_positive(variance)
+            quantized_moment, moment_scale = _quantize_dynamic(moment, signed=True)
+            quantized_variance, variance_scale = _quantize_dynamic(variance, signed=False)
             block_moment_q.copy_(quantized_moment)
             block_variance_q.copy_(quantized_variance)
             moment_scales[block_index].copy_(moment_scale)
             variance_scales[block_index].copy_(variance_scale)
 
 
-def _quantize_symmetric(tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-    """Quantize a signed FP32 tensor to uint8 with one bucket-level scale."""
+def set_optimizer_state_precision(
+    parameter: torch.nn.Parameter,
+    precision: str,
+    *,
+    role: str | None = None,
+) -> torch.nn.Parameter:
+    """Attach an explicit AdamW8bit state policy to a parameter.
+
+    This is intentionally parameter metadata rather than name matching, so it
+    survives bucket construction and works for non-standard model layouts.
+    """
+
+    normalized = _normalize_state_precision(precision)
+    parameter._areno_optimizer_state_precision = normalized
+    if role is not None:
+        parameter._areno_optimizer_role = str(role)
+    return parameter
+
+
+def _normalize_parameter_policies(
+    params: Iterable[torch.nn.Parameter] | Iterable[Mapping[str, Any]],
+) -> tuple[list[torch.nn.Parameter], dict[int, str], dict[int, str]]:
+    values = list(params)
+    flattened: list[torch.nn.Parameter] = []
+    precision_by_id: dict[int, str] = {}
+    role_by_id: dict[int, str] = {}
+    seen: set[int] = set()
+    for value in values:
+        if isinstance(value, Mapping):
+            group_params = list(value.get("params", ()))
+            group_precision = value.get("state_precision")
+            group_role = value.get("role")
+        else:
+            group_params = [value]
+            group_precision = None
+            group_role = None
+        for parameter in group_params:
+            if not isinstance(parameter, torch.nn.Parameter):
+                raise TypeError("AdamW8bit params must contain torch.nn.Parameter values")
+            identity = id(parameter)
+            explicit = getattr(parameter, "_areno_optimizer_state_precision", group_precision)
+            if explicit is not None:
+                precision_by_id[identity] = _normalize_state_precision(str(explicit))
+            role = getattr(parameter, "_areno_optimizer_role", group_role)
+            if role is not None:
+                role_by_id[identity] = str(role)
+            if identity not in seen:
+                flattened.append(parameter)
+                seen.add(identity)
+    return flattened, precision_by_id, role_by_id
+
+
+def _normalize_state_precision(precision: str) -> str:
+    aliases = {"uint8": "8bit", "int8": "8bit", "float32": "fp32"}
+    normalized = aliases.get(precision.lower(), precision.lower())
+    if normalized not in _VALID_STATE_PRECISIONS:
+        raise ValueError(f"state_precision must be one of {sorted(_VALID_STATE_PRECISIONS)}, got {precision!r}")
+    return normalized
+
+
+def _dynamic_codebook(device: torch.device, *, signed: bool) -> torch.Tensor:
+    key = (device, signed)
+    codebook = _CODEBOOK_CACHE.get(key)
+    if codebook is None:
+        values = SIGNED_DYNAMIC_MAP if signed else UNSIGNED_DYNAMIC_MAP
+        codebook = torch.tensor(values, device=device, dtype=torch.float32)
+        _CODEBOOK_CACHE[key] = codebook
+    return codebook
+
+
+def _quantize_dynamic(tensor: torch.Tensor, *, signed: bool) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize one FP32 block with the paper's dynamic-tree codebook."""
 
     if tensor.numel() == 0:
-        return tensor.to(dtype=torch.uint8), torch.ones((), device=tensor.device, dtype=torch.float32)
-    scale = tensor.abs().amax().div(127.0).clamp_min(1.0e-30)
-    quantized = torch.clamp(torch.round(tensor / scale) + 128.0, 0.0, 255.0).to(dtype=torch.uint8)
-    return quantized, scale.to(dtype=torch.float32)
+        return tensor.to(dtype=torch.uint8), torch.zeros((), device=tensor.device, dtype=torch.float32)
+    scale = (tensor.abs().amax() if signed else tensor.clamp_min(0).amax()).to(dtype=torch.float32)
+    normalized = tensor.float().div(scale.clamp_min(torch.finfo(torch.float32).tiny))
+    if not signed:
+        normalized.clamp_min_(0.0)
+    codebook = _dynamic_codebook(tensor.device, signed=signed)
+    boundaries = (codebook[:-1] + codebook[1:]) * 0.5
+    codes = torch.bucketize(normalized, boundaries).to(dtype=torch.uint8)
+    return codes, scale
 
 
-def _dequantize_symmetric(quantized: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
-    """Dequantize signed uint8 moments back to FP32."""
+def _dequantize_dynamic(quantized: torch.Tensor, scale: torch.Tensor, *, signed: bool) -> torch.Tensor:
+    codebook = _dynamic_codebook(quantized.device, signed=signed)
+    return codebook[quantized.long()].mul_(scale)
 
-    return (quantized.to(dtype=torch.float32) - 128.0).mul_(scale)
 
+# Private convenience aliases used by focused quantization tests.
+def _quantize_symmetric(tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    return _quantize_dynamic(tensor, signed=True)
 
-def _quantize_positive(tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-    """Quantize one non-negative FP32 block to uint8."""
 
-    if tensor.numel() == 0:
-        return tensor.to(dtype=torch.uint8), torch.ones((), device=tensor.device, dtype=torch.float32)
-    scale = tensor.amax().div(255.0).clamp_min(1.0e-30)
-    quantized = torch.clamp(torch.round(tensor / scale), 0.0, 255.0).to(dtype=torch.uint8)
-    return quantized, scale.to(dtype=torch.float32)
+def _dequantize_symmetric(quantized: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    return _dequantize_dynamic(quantized, scale, signed=True)
+
+
+def _quantize_positive(tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    return _quantize_dynamic(tensor, signed=False)
 
 
 def _dequantize_positive(quantized: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
-    """Dequantize non-negative uint8 moments back to FP32."""
+    return _dequantize_dynamic(quantized, scale, signed=False)
+
 
-    return quantized.to(dtype=torch.float32).mul_(scale)
+def _load_optional_state_tensor(
+    saved: Any,
+    bucket: _MasterBucket,
+    device: torch.device,
+) -> torch.Tensor | None:
+    if saved is None:
+        return None
+    value = saved.detach().to(device=device, dtype=torch.float32).view(-1)
+    if value.numel() != bucket.shard_numel:
+        raise ValueError(
+            f"AdamW8bit checkpoint has {value.numel()} FP32 state values for a bucket requiring {bucket.shard_numel}"
+        )
+    return value.clone()
+
+
+def _load_optional_quantized_tensor(
+    saved: Any,
+    bucket: _MasterBucket,
+    device: torch.device,
+) -> torch.Tensor | None:
+    if saved is None:
+        return None
+    value = saved.detach().to(device=device, dtype=torch.uint8).view(-1)
+    if value.numel() != bucket.shard_numel:
+        raise ValueError(
+            f"AdamW8bit checkpoint has {value.numel()} quantized state values "
+            f"for a bucket requiring {bucket.shard_numel}"
+        )
+    return value.clone()
 
 
 def _cpu_clone(value: torch.Tensor | None) -> torch.Tensor | None:
diff --git a/areno/engine/optim/dynamic_quant.py b/areno/engine/optim/dynamic_quant.py
new file mode 100644
index 00000000..2874072c
--- /dev/null
+++ b/areno/engine/optim/dynamic_quant.py
@@ -0,0 +1,61 @@
+"""Reference codebooks for Dettmers et al. block-wise dynamic quantization."""
+
+from __future__ import annotations
+
+
+def create_dynamic_map(*, signed: bool, max_exponent_bits: int = 7, total_bits: int = 8) -> tuple[float, ...]:
+    """Build the paper's sorted dynamic-tree codebook without a torch dependency.
+
+    The construction matches the reference implementation released with
+    ``8-Bit Optimizers via Block-wise Quantization``. Signed maps reserve one
+    sign bit; unsigned maps reclaim it for an additional fractional bit.
+    """
+
+    if total_bits != 8:
+        raise ValueError("AReno dynamic optimizer states currently require total_bits=8")
+    if max_exponent_bits < 1 or max_exponent_bits >= total_bits:
+        raise ValueError("max_exponent_bits must be between 1 and total_bits - 1")
+
+    values: list[float] = []
+    non_sign_bits = total_bits - 1
+    additional_items = 2 ** (non_sign_bits - max_exponent_bits) - 1
+    last_index = 0
+    for index in range(max_exponent_bits):
+        last_index = index
+        fraction_items = 2 ** (index + non_sign_bits - max_exponent_bits + (0 if signed else 1)) + 1
+        step = 0.9 / (fraction_items - 1)
+        means = (0.1 + (item + 0.5) * step for item in range(fraction_items - 1))
+        scale = 10 ** (-(max_exponent_bits - 1) + index)
+        positive = [scale * mean for mean in means]
+        values.extend(positive)
+        if signed:
+            values.extend(-value for value in positive)
+
+    if additional_items > 0:
+        step = 0.9 / additional_items
+        means = (0.1 + (item + 0.5) * step for item in range(additional_items))
+        scale = 10 ** (-(max_exponent_bits - 1) + last_index)
+        positive = [scale * mean for mean in means]
+        values.extend(positive)
+        if signed:
+            values.extend(-value for value in positive)
+
+    values.extend((0.0, 1.0))
+    if len(values) != 2**total_bits:
+        raise AssertionError(f"dynamic codebook has {len(values)} entries, expected {2**total_bits}")
+    return tuple(sorted(values))
+
+
+SIGNED_DYNAMIC_MAP = create_dynamic_map(signed=True)
+UNSIGNED_DYNAMIC_MAP = create_dynamic_map(signed=False)
+SIGNED_DYNAMIC_ZERO = SIGNED_DYNAMIC_MAP.index(0.0)
+UNSIGNED_DYNAMIC_ZERO = UNSIGNED_DYNAMIC_MAP.index(0.0)
+
+
+__all__ = [
+    "SIGNED_DYNAMIC_MAP",
+    "SIGNED_DYNAMIC_ZERO",
+    "UNSIGNED_DYNAMIC_MAP",
+    "UNSIGNED_DYNAMIC_ZERO",
+    "create_dynamic_map",
+]
diff --git a/areno/engine/training.py b/areno/engine/training.py
index e0dcf585..3aac40af 100644
--- a/areno/engine/training.py
+++ b/areno/engine/training.py
@@ -163,6 +163,7 @@ def _train_step(
         grad_norm = None
         multimodal_grad_metrics = None
         clipped_grad_norm = None
+        optimizer_state_metrics = None
         if stepped:
             self._sync_data_parallel_gradients()
             self._sync_tensor_parallel_replicated_gradients()
@@ -185,6 +186,12 @@ def _train_step(
             worker.optimizer.lr = current_lr
             multimodal_lrs = self._set_multimodal_lrs(worker._global_step + 1)
             worker.optimizer.step()
+            state_memory_metrics = getattr(worker.optimizer, "state_memory_metrics", None)
+            if (
+                callable(state_memory_metrics)
+                and getattr(worker.optimizer, "state_quantizer", None) == "dynamic-tree-v1"
+            ):
+                optimizer_state_metrics = {f"adam8_{name}": value for name, value in state_memory_metrics().items()}
             worker.optimizer.zero_grad(set_to_none=True)
             worker._global_step += 1
             if worker.adapter_registry is not None:
@@ -216,6 +223,7 @@ def _train_step(
                     {"grad_norm": grad_norm} if grad_norm is not None else None,
                     multimodal_grad_metrics,
                     {"clipped_grad_norm": clipped_grad_norm} if clipped_grad_norm is not None else None,
+                    optimizer_state_metrics,
                 ),
             }
         return None
diff --git a/docs/_static/kaggle-dataset-mixing/kaggle-01-source-revision.png b/docs/_static/kaggle-dataset-mixing/kaggle-01-source-revision.png
new file mode 100644
index 00000000..123c9631
Binary files /dev/null and b/docs/_static/kaggle-dataset-mixing/kaggle-01-source-revision.png differ
diff --git a/docs/_static/kaggle-dataset-mixing/kaggle-02-environment.png b/docs/_static/kaggle-dataset-mixing/kaggle-02-environment.png
new file mode 100644
index 00000000..5f8de7e0
Binary files /dev/null and b/docs/_static/kaggle-dataset-mixing/kaggle-02-environment.png differ
diff --git a/docs/_static/kaggle-dataset-mixing/kaggle-03-model-preparation.png b/docs/_static/kaggle-dataset-mixing/kaggle-03-model-preparation.png
new file mode 100644
index 00000000..69207489
Binary files /dev/null and b/docs/_static/kaggle-dataset-mixing/kaggle-03-model-preparation.png differ
diff --git a/docs/_static/kaggle-dataset-mixing/kaggle-04-fp16-checkpoint.png b/docs/_static/kaggle-dataset-mixing/kaggle-04-fp16-checkpoint.png
new file mode 100644
index 00000000..82db56a0
Binary files /dev/null and b/docs/_static/kaggle-dataset-mixing/kaggle-04-fp16-checkpoint.png differ
diff --git a/docs/_static/kaggle-dataset-mixing/kaggle-05-token-limit-128-config.png b/docs/_static/kaggle-dataset-mixing/kaggle-05-token-limit-128-config.png
new file mode 100644
index 00000000..77252c45
Binary files /dev/null and b/docs/_static/kaggle-dataset-mixing/kaggle-05-token-limit-128-config.png differ
diff --git a/docs/_static/kaggle-dataset-mixing/kaggle-06-token-limit-128-result.png b/docs/_static/kaggle-dataset-mixing/kaggle-06-token-limit-128-result.png
new file mode 100644
index 00000000..0ba8a540
Binary files /dev/null and b/docs/_static/kaggle-dataset-mixing/kaggle-06-token-limit-128-result.png differ
diff --git a/docs/_static/kaggle-dataset-mixing/kaggle-07-structured-plan.png b/docs/_static/kaggle-dataset-mixing/kaggle-07-structured-plan.png
new file mode 100644
index 00000000..b6a27dae
Binary files /dev/null and b/docs/_static/kaggle-dataset-mixing/kaggle-07-structured-plan.png differ
diff --git a/docs/_static/kaggle-dataset-mixing/kaggle-08-token-limit-256-config.png b/docs/_static/kaggle-dataset-mixing/kaggle-08-token-limit-256-config.png
new file mode 100644
index 00000000..52c149ab
Binary files /dev/null and b/docs/_static/kaggle-dataset-mixing/kaggle-08-token-limit-256-config.png differ
diff --git a/docs/_static/kaggle-dataset-mixing/kaggle-09-token-limit-256-result.png b/docs/_static/kaggle-dataset-mixing/kaggle-09-token-limit-256-result.png
new file mode 100644
index 00000000..ebe283f3
Binary files /dev/null and b/docs/_static/kaggle-dataset-mixing/kaggle-09-token-limit-256-result.png differ
diff --git a/docs/_static/kaggle-dataset-mixing/kaggle-10-500-step-config.png b/docs/_static/kaggle-dataset-mixing/kaggle-10-500-step-config.png
new file mode 100644
index 00000000..fe23dd7b
Binary files /dev/null and b/docs/_static/kaggle-dataset-mixing/kaggle-10-500-step-config.png differ
diff --git a/docs/_static/kaggle-dataset-mixing/kaggle-11-500-step-result.png b/docs/_static/kaggle-dataset-mixing/kaggle-11-500-step-result.png
new file mode 100644
index 00000000..ea7ced1e
Binary files /dev/null and b/docs/_static/kaggle-dataset-mixing/kaggle-11-500-step-result.png differ
diff --git a/docs/_static/kaggle-dataset-mixing/kaggle-12-500-step-summary.png b/docs/_static/kaggle-dataset-mixing/kaggle-12-500-step-summary.png
new file mode 100644
index 00000000..be6ed9e0
Binary files /dev/null and b/docs/_static/kaggle-dataset-mixing/kaggle-12-500-step-summary.png differ
diff --git a/docs/cli/training.rst b/docs/cli/training.rst
index 5034c984..350bdbcb 100644
--- a/docs/cli/training.rst
+++ b/docs/cli/training.rst
@@ -71,6 +71,112 @@ trainer schema. Without a loader, Areno passes dataset rows through unchanged,
 except for SFT where a loader is required. See :doc:`dataset_loaders` for the
 per-algorithm loader contracts.
 
+``--dataset-mix-config PATH``
+   JSON manifest for deterministic weighted SFT dataset mixing. This option is
+   mutually exclusive with ``--dataset-path`` and ``--dataset-source``.
+
+``--dataset-source NAME=PATH:WEIGHT``
+   Command-line shorthand for a weighted SFT source. Repeat the option at least
+   twice. It is mutually exclusive with ``--dataset-path`` and
+   ``--dataset-mix-config``:
+
+   .. code-block:: bash
+
+      areno train \
+        --algo sft \
+        --ckpt Qwen/Qwen3-0.6B \
+        --model-hub hf \
+        --dataset-source alpaca_cleaned=yahma/alpaca-cleaned:0.6 \
+        --dataset-source stanford_alpaca=tatsu-lab/alpaca:0.3 \
+        --dataset-source alpaca_gpt4=vicgalle/alpaca-gpt4:0.1 \
+        --dataset-mix-samples-per-epoch 10000 \
+        --dataset-loader-fn examples/sft/alpaca/dataset_loader.py
+
+   Any number of sources (two or more) can be repeated. Weights are normalized
+   across all sources and mean probabilities of selecting a **row**, not a
+   token or a batch. The shorthand uses seed ``42``, ``cycle`` exhaustion, and
+   source shuffling. Override these with ``--dataset-mix-seed``,
+   ``--dataset-mix-exhaustion``, and
+   ``--dataset-mix-samples-per-epoch``.
+
+``--dataset-mix-samples-per-epoch INTEGER``
+   Fixed row-sampling budget for a command-line ``cycle`` mix. If omitted,
+   AReno uses the sum of the loaded source row counts. A fixed budget keeps
+   many-source training bounded while preserving the requested distribution;
+   small sources repeat as necessary.
+
+``--dataset-mix-exhaustion [cycle|stop|renormalize]``
+   Exhaustion policy for command-line sources. ``cycle`` is the default because
+   it keeps weights meaningful for the whole fixed budget. ``stop`` and
+   ``renormalize`` do not accept a sample budget.
+
+``--dataset-mix-seed INTEGER``
+   Seed for source selection and per-source shuffling. Default: ``42``.
+
+Both mixing forms currently support map-style datasets only, and the shared
+``--dataset-loader-fn`` is applied independently to each source.
+
+The manifest requires ``version: 1``, an integer ``seed``, an explicit
+``exhaustion`` policy, and at least two uniquely named sources with finite,
+positive weights. Unknown version-1 manifest or source fields are rejected so
+misspelled settings cannot silently fall back to defaults. Weights are
+normalized automatically:
+
+.. code-block:: json
+
+   {
+     "version": 1,
+     "seed": 42,
+     "exhaustion": "cycle",
+     "shuffle_within_sources": true,
+     "samples_per_epoch": 10000,
+     "sources": [
+       {"name": "math", "path": "math.jsonl", "weight": 0.7},
+       {"name": "code", "path": "code.jsonl", "weight": 0.3}
+     ]
+   }
+
+Existing relative local paths, explicit ``./`` or ``../`` paths, and supported
+local file suffixes are resolved from the manifest directory. Other values are
+handled as remote dataset references using ``--model-hub``. The exhaustion
+policies have deliberately different consumption behavior:
+
+* ``stop`` ends when one selected source is exhausted. Records are not
+  repeated, but other sources may remain partially unused.
+* ``cycle`` restarts exhausted sources and emits exactly
+  ``samples_per_epoch`` rows. When the field/flag is omitted, both manifest and
+  command-line forms resolve the budget to the sum of source row counts.
+* ``renormalize`` removes an exhausted source and renormalizes the remaining
+  weights. Every source record is emitted exactly once.
+
+For ``K`` sources and budget ``N``, source ``i`` receives about
+``N * normalized_weight_i`` rows. This is stochastic: a source whose expected
+count is below one can receive zero rows in an epoch, so AReno emits a plan
+warning. Use a larger budget when every low-weight source needs practical
+coverage.
+
+The seed controls source selection and per-source ordering. AReno derives a
+different deterministic ordering for each epoch. Source identity is attached
+under the reserved ``__areno_meta__`` field. The run prints a sample-free
+summary, logs the plan for each epoch, and records both a mix-spec hash and
+schedule hash. These hashes identify sampler configuration and index ordering
+only; they do not prove that a mutable remote dataset still has identical
+contents.
+The cumulative ``stage=dataset_mix_progress`` event reports scheduled,
+filtered, and successfully trained rows plus trained target-token counts and
+both row/token proportions. Each epoch plan is also written as
+``dataset_mix_plan.<pid>.epoch-<n>.json`` to ``--metrics-log-dir``.
+
+V1 guarantees deterministic replay from the beginning of an epoch. AReno
+checkpoints currently do not persist an optimizer state plus mid-epoch dataset
+cursor, so exact mid-epoch resume is outside this contract.
+
+The map-style implementation precomputes an index schedule for the current
+epoch. Very large mixes therefore require memory proportional to the planned
+row count; streaming datasets are not supported in this first version.
+See ``examples/sft/mixed`` for a copyable local example and an invalid-weight
+boundary case.
+
 ``--algo TEXT``
    Training algorithm registered in ``areno.api``. Default: ``gspo``.
 
@@ -380,8 +486,22 @@ in its description; flags for other algorithms are ignored.
    Policy optimizer Adam beta2. Default: ``0.999``.
 
 ``--adam-8bit``
-   Use 8-bit Adam moment states instead of FP32 Adam states. Supported by both
-   native backends; validate convergence when changing optimizer precision.
+   Use block-wise 8-bit Adam moment states instead of FP32 Adam states.
+   CUDA and MLX use the signed dynamic-tree codebook for the first moment and
+   the unsigned dynamic codebook for the second moment from *8-Bit Optimizers
+   via Block-wise Quantization*. Token-embedding weights and gradients retain
+   their normal model precision, while their optimizer moments remain FP32 to
+   avoid quantizing embedding-gradient outliers. Other initialized moment
+   state uses two bytes per parameter plus FP32 block scales.
+   CUDA training metrics expose ``adam8_quantized_state_bytes``,
+   ``adam8_fp32_exempt_bytes``, ``adam8_block_metadata_bytes``, and
+   ``adam8_total_bytes`` for initialized DP-local optimizer state.
+
+   This option does not insert a normalization layer or reinitialize a loaded
+   embedding. The paper's Stable Embedding forward architecture is not enabled
+   for AReno's current RoPE-only language models because they do not expose the
+   compatible additive-position-embedding boundary. Supported by both native
+   backends; validate convergence when changing optimizer precision.
 
 ``--adam-4bit``
    Use packed block-wise 4-bit Adam moment states. This option is CUDA-only
diff --git a/docs/cookbook/kaggle-dataset-mixing-validation.rst b/docs/cookbook/kaggle-dataset-mixing-validation.rst
new file mode 100644
index 00000000..ac334561
--- /dev/null
+++ b/docs/cookbook/kaggle-dataset-mixing-validation.rst
@@ -0,0 +1,790 @@
+Kaggle dataset-mixing validation
+================================
+
+This guide records the reproducible Kaggle GPU validation for deterministic
+weighted SFT dataset mixing introduced by Issue #198. It covers environment
+setup, model preparation, a three-source training run, deterministic-plan
+verification, filtering analysis, and the evidence that should be captured
+from the Kaggle notebook.
+
+The validation targets the data-mixing and SFT execution contracts. It is not
+an evaluation of downstream model quality. Prompts, responses, credentials,
+and Hugging Face tokens must not be included in screenshots or attached logs.
+
+Validation summary
+------------------
+
+The tested implementation completed all of the following:
+
+* loaded and normalized three public Alpaca-contract datasets;
+* produced the requested 60/30/10 deterministic sample schedule;
+* reproduced the same schedule hash across three independent runs;
+* completed 500 SFT optimizer steps with two-GPU tensor parallelism;
+* trained 1,000 accepted rows and 77,533 target tokens;
+* exposed source-specific scheduled, filtered, trained, and token counts;
+* completed without NaN, Inf, CUDA, or worker errors.
+
+The implementation under test was commit ``e33e40d``. Later commits on the PR
+branch only add or refine documentation.
+
+Validated environment
+---------------------
+
+The run was captured from a fresh Kaggle notebook with Internet access and the
+GPU accelerator enabled.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 32 68
+
+   * - Component
+     - Observed value
+   * - Operating system
+     - Linux 6.12, x86_64
+   * - Python
+     - 3.12.13
+   * - PyTorch
+     - 2.10.0+cu128
+   * - PyTorch CUDA build
+     - 12.8
+   * - NVCC
+     - 12.8 (V12.8.93)
+   * - Accelerator
+     - 2 x Tesla T4; the SFT runs used both GPUs with tensor parallelism
+   * - Per-GPU memory
+     - 15,360 MiB reported by ``nvidia-smi``
+   * - GPU compute capability
+     - 7.5
+   * - Attention backend
+     - ``native``
+   * - Model
+     - Qwen3-0.6B converted to FP16
+
+.. figure:: ../_static/kaggle-dataset-mixing/kaggle-02-environment.png
+   :alt: Kaggle PyTorch, CUDA, Tesla T4, and areno_accel environment check
+   :width: 100%
+
+   Kaggle environment check showing PyTorch 2.10.0+cu128, CUDA 12.8, a
+   visible Tesla T4, and a successful import of the compiled
+   ``areno_accel`` extension.
+
+1. Check out the PR branch
+--------------------------
+
+Clone the fork branch into the ephemeral Kaggle working directory:
+
+.. code-block:: bash
+
+   cd /kaggle/working
+   git clone \
+     --branch codex/issue-198-dataset-mixing \
+     --single-branch \
+     https://github.com/lkxdsb/AReno.git
+   cd /kaggle/working/AReno
+
+Record the exact source revision before installing:
+
+.. code-block:: bash
+
+   git branch --show-current
+   git rev-parse HEAD
+   git status --short --branch
+
+The branch should be ``codex/issue-198-dataset-mixing`` and the working tree
+should be clean. The exact commit can be newer than the validated
+implementation commit when it contains documentation-only updates.
+
+.. figure:: ../_static/kaggle-dataset-mixing/kaggle-01-source-revision.png
+   :alt: Kaggle checkout of the Issue 198 branch and exact source revision
+   :width: 100%
+
+   The notebook cloned ``codex/issue-198-dataset-mixing`` and recorded
+   implementation revision
+   ``e33e40de3419e37a44bf5279374ada6d6c37eae3`` before installation.
+
+2. Install AReno and compile the CUDA extension
+-----------------------------------------------
+
+Kaggle's T4 uses compute capability 7.5. Restricting the extension build to
+``sm_75`` reduces build time and temporary disk use:
+
+.. code-block:: bash
+
+   export CUDA_HOME=/usr/local/cuda
+   export TORCH_CUDA_ARCH_LIST=7.5
+   export MAX_JOBS=2
+
+   python -m pip install "setuptools>=69" wheel psutil ninja
+   python -m pip install "flash-linear-attention>=0.2"
+   python -m pip install -e . --no-build-isolation
+
+Do not set ``ARENO_BUILD_EXT=0`` for this validation. Native attention still
+uses the compiled ``areno_accel`` extension for training operations.
+
+Run the built-in diagnostic after installation:
+
+.. code-block:: bash
+
+   areno check
+
+The check must see CUDA, the model/runtime dependencies, and the compiled
+extension. A successful package installation without ``areno_accel`` is not a
+valid training environment. A missing ``flash_attn`` warning is acceptable for
+this test because the command explicitly selects ``--attn-backend native``;
+the diagnostic must still exit with code zero.
+
+The environment figure above is the retained installation evidence: it proves
+that CUDA is available and the compiled extension imports successfully. For a
+new reproduction, retain the complete ``areno check`` output as well. The
+shell here-document warning visible in the captured cell came from the
+notebook wrapper and did not affect the successful CUDA or extension imports.
+
+3. Prepare the FP16 checkpoint
+------------------------------
+
+The T4 run used a local FP16 copy of Qwen3-0.6B:
+
+.. code-block:: python
+
+   import torch
+   from transformers import AutoModelForCausalLM, AutoTokenizer
+
+   model_id = "Qwen/Qwen3-0.6B"
+   output_dir = "/kaggle/working/qwen3-0.6b-fp16"
+
+   tokenizer = AutoTokenizer.from_pretrained(model_id)
+   model = AutoModelForCausalLM.from_pretrained(
+       model_id,
+       dtype=torch.float16,
+       low_cpu_mem_usage=True,
+   )
+   model.save_pretrained(output_dir, safe_serialization=True)
+   tokenizer.save_pretrained(output_dir)
+
+   print("checkpoint:", output_dir)
+   print("dtype:", next(model.parameters()).dtype)
+
+Warnings about anonymous Hugging Face rate limits or tied-weight metadata do
+not invalidate the conversion. The decisive evidence is a completed
+``model.safetensors`` write and ``torch.float16`` parameter dtype.
+
+The captured notebook used the older ``torch_dtype=torch.float16`` spelling.
+The reproduction cell above uses ``dtype=torch.float16`` because current
+Transformers versions deprecate ``torch_dtype``.
+
+.. figure:: ../_static/kaggle-dataset-mixing/kaggle-03-model-preparation.png
+   :alt: Kaggle code cell that downloads and converts Qwen3 0.6B to FP16
+   :width: 100%
+
+   Model-preparation cell used to download Qwen3-0.6B and write a local FP16
+   checkpoint.
+
+.. figure:: ../_static/kaggle-dataset-mixing/kaggle-04-fp16-checkpoint.png
+   :alt: Completed Qwen3 0.6B download and FP16 checkpoint write
+   :width: 100%
+
+   Completed weight download and checkpoint write. The final output confirms
+   ``/kaggle/working/qwen3-0.6b-fp16`` and ``torch.float16``.
+
+4. Define the three-source mix
+------------------------------
+
+The GPU validation used three datasets accepted by the shared Alpaca loader:
+
+.. list-table::
+   :header-rows: 1
+   :widths: 24 48 14 14
+
+   * - Source name
+     - Dataset reference
+     - Weight
+     - Rows available
+   * - ``alpaca``
+     - ``yahma/alpaca-cleaned``
+     - 0.6
+     - 51,760
+   * - ``stanford``
+     - ``tatsu-lab/alpaca``
+     - 0.3
+     - 52,002
+   * - ``gpt4``
+     - ``vicgalle/alpaca-gpt4``
+     - 0.1
+     - 52,002
+
+All three sources are normalized by
+``examples/sft/alpaca/dataset_loader.py`` into the SFT ``prompt`` /
+``response`` contract. The source names are diagnostic identifiers and do not
+need to match repository names.
+
+Inline source weights are normalized automatically. The run used the inline
+defaults:
+
+* seed: ``42``;
+* exhaustion: ``cycle``;
+* shuffle within sources: enabled;
+* weight unit: ``sample``;
+* samples per epoch: sum of loaded source row counts when omitted.
+
+5. Run the 500-step validation
+------------------------------
+
+The notebook first ran 50-step checks at 128 and 256 tokens, then ran the final
+500-step validation at 256 tokens. Sections 8 through 10 retain the commands
+and final output from all three passes.
+
+The following is the command corresponding to the recorded 500-step result:
+
+.. code-block:: bash
+
+   areno train \
+     --algo sft \
+     --ckpt /kaggle/working/qwen3-0.6b-fp16 \
+     --model-hub hf \
+     --dataset-source alpaca=yahma/alpaca-cleaned:0.6 \
+     --dataset-source stanford=tatsu-lab/alpaca:0.3 \
+     --dataset-source gpt4=vicgalle/alpaca-gpt4:0.1 \
+     --dataset-loader-fn examples/sft/alpaca/dataset_loader.py \
+     --world-size 2 --tp-size 2 \
+     --batch-size 2 --mini-bs 1 \
+     --max-steps 500 \
+     --max-prompt-tokens 256 --max-new-tokens 256 \
+     --attn-backend native
+
+This run intentionally omitted ``--save-path`` because its purpose was feature
+validation rather than preserving a fine-tuned model. Add a save path for a
+training run whose weights must survive process exit:
+
+.. code-block:: bash
+
+   --save-path /kaggle/working/qwen3-mixed-sft \
+   --save-interval 500
+
+The recorded plan used the automatic 155,764-row epoch budget. For a shorter
+bounded plan, add:
+
+.. code-block:: bash
+
+   --dataset-mix-samples-per-epoch 5000
+
+This option bounds scheduled rows. It does not limit source download/loading,
+and it does not guarantee 5,000 accepted training rows. In this configuration,
+``500 steps x batch size 2`` requires 1,000 accepted rows; filtering causes the
+trainer to scan more than 1,000 scheduled rows to fill those batches.
+
+.. figure:: ../_static/kaggle-dataset-mixing/kaggle-10-500-step-config.png
+   :alt: Kaggle 500-step three-source SFT command and resolved configuration
+   :width: 100%
+
+   The 500-step command and the beginning of AReno's resolved configuration.
+   It records three inline sources, ``world_size=2``, ``tp_size=2``, batch size
+   two, 256-token limits, and the native attention backend.
+
+6. Inspect the structured plan
+------------------------------
+
+When no explicit metrics directory is supplied, the default directory is
+``/tmp/areno/tfevent``. Read the newest sample-free plan:
+
+.. code-block:: python
+
+   import glob
+   import json
+   import os
+
+   plan_path = max(
+       glob.glob("/tmp/areno/tfevent/dataset_mix_plan.*.json"),
+       key=os.path.getmtime,
+   )
+   with open(plan_path, encoding="utf-8") as handle:
+       plan = json.load(handle)
+
+   print("mix_spec_hash:", plan["mix_spec_hash"])
+   print("schedule_hash:", plan["schedule_hash"])
+   print("planned_rows:", plan["planned_rows"])
+   for source in plan["sources"]:
+       print(
+           source["name"],
+           "requested=", source["weight_requested"],
+           "selected=", source["rows_selected"],
+           "observed=", source["observed_proportion"],
+           "duplicates=", source["duplicates"],
+       )
+
+The recorded epoch-zero plan was:
+
+.. list-table::
+   :header-rows: 1
+   :widths: 20 18 18 18 18
+
+   * - Source
+     - Requested
+     - Selected
+     - Observed
+     - Duplicates
+   * - ``alpaca``
+     - 60%
+     - 93,486
+     - 60.0177%
+     - 41,726
+   * - ``stanford``
+     - 30%
+     - 46,775
+     - 30.0294%
+     - 0
+   * - ``gpt4``
+     - 10%
+     - 15,503
+     - 9.9529%
+     - 0
+
+The full-epoch ``cycle`` plan repeats the smaller high-weight Alpaca source
+after its 51,760 unique rows are exhausted. The 500-step run consumed only the
+first 1,171 scheduled rows, so this full-plan duplicate count does not mean
+that the short validation consumed 41,726 duplicates.
+
+.. figure:: ../_static/kaggle-dataset-mixing/kaggle-07-structured-plan.png
+   :alt: Structured dataset-mixing plan with schedule hash and source counts
+   :width: 100%
+
+   Structured plan inspection showing the schedule hash, 155,764 planned
+   rows, and requested, planned, and selected values for all three sources.
+
+7. Verify deterministic replay
+------------------------------
+
+The 128-token run, 256-token run, and 500-step run all used the same source
+contract and produced:
+
+.. code-block:: text
+
+   mix_spec_hash:
+   sha256:fda6d7cbdab898b0ddafb6e6f37fa6c81b581da73b4bde7ebdfda8a0438057b4
+
+   schedule_hash:
+   sha256:26e61327f09b9c27a0e7219a40e2d17bebc5604571b47a1b8a3dc4119aad420b
+
+Token limits do not participate in source scheduling, so changing 128 to 256
+tokens must not change either hash. Compare all saved plans:
+
+.. code-block:: python
+
+   import glob
+   import json
+   import os
+
+   for path in sorted(
+       glob.glob("/tmp/areno/tfevent/dataset_mix_plan.*.json"),
+       key=os.path.getmtime,
+   ):
+       with open(path, encoding="utf-8") as handle:
+           plan = json.load(handle)
+       print(os.path.basename(path), plan["schedule_hash"])
+
+The structured-plan figure and the final 500-step summary below both report
+``sha256:26e61327f09b9c27a0e7219a40e2d17bebc5604571b47a1b8a3dc4119aad420b``.
+Together with the matching plan artifacts observed in the intervening
+256-token run,
+this verifies that token limits and step count did not alter the schedule.
+
+8. Compare token-limit filtering
+--------------------------------
+
+Before the final run, two 50-step runs compared 128-token and 256-token
+prompt/response budgets. They used the command from section 5 with
+``--max-steps 50`` and, respectively:
+
+.. code-block:: bash
+
+   --max-prompt-tokens 128 --max-new-tokens 128
+
+.. code-block:: bash
+
+   --max-prompt-tokens 256 --max-new-tokens 256
+
+.. list-table::
+   :header-rows: 1
+   :widths: 32 22 22
+
+   * - Metric
+     - 128 / 128
+     - 256 / 256
+   * - Scheduled rows
+     - 165
+     - 117
+   * - Filtered rows
+     - 65
+     - 17
+   * - Filter rate
+     - 39.39%
+     - 14.53%
+   * - Trained rows
+     - 100
+     - 100
+   * - Target tokens trained
+     - 4,455
+     - 8,295
+   * - Target tokens per trained row
+     - 44.55
+     - 82.95
+
+The 256-token configuration reduced filtering by 24.86 percentage points and
+nearly doubled the supervised target-token volume without destabilizing the
+two-GPU tensor-parallel run. It was therefore selected for the 500-step
+validation.
+
+.. figure:: ../_static/kaggle-dataset-mixing/kaggle-05-token-limit-128-config.png
+   :alt: Kaggle 128-token 50-step SFT command and resolved configuration
+   :width: 100%
+
+   First validation command: 50 steps with 128 prompt tokens and 128 response
+   tokens. The resolved input summary confirms that three command-line sources
+   were selected.
+
+.. figure:: ../_static/kaggle-dataset-mixing/kaggle-06-token-limit-128-result.png
+   :alt: Final progress and completion logs from the 128-token run
+   :width: 100%
+
+   Final 128-token progress. The run reached 50 steps after scheduling 165
+   rows, filtering 65, and training 100.
+
+.. figure:: ../_static/kaggle-dataset-mixing/kaggle-08-token-limit-256-config.png
+   :alt: Kaggle 256-token 50-step SFT command and resolved configuration
+   :width: 100%
+
+   Second validation command: the same 50-step run with both token limits
+   raised to 256. All dataset-source and runtime settings remain unchanged.
+
+.. figure:: ../_static/kaggle-dataset-mixing/kaggle-09-token-limit-256-result.png
+   :alt: Final progress and completion logs from the 256-token run
+   :width: 100%
+
+   Final 256-token progress. The run reached 50 steps after scheduling 117
+   rows, filtering 17, and training 100.
+
+9. Inspect the 500-step result
+------------------------------
+
+The final progress event reported:
+
+.. list-table::
+   :header-rows: 1
+   :widths: 34 22
+
+   * - Metric
+     - Result
+   * - Optimizer steps
+     - 500
+   * - Scheduled rows consumed
+     - 1,171
+   * - Filtered rows
+     - 171
+   * - Overall filter rate
+     - 14.60%
+   * - Trained rows
+     - 1,000
+   * - Target tokens trained
+     - 77,533
+   * - Mean target tokens per trained row
+     - 77.533
+   * - Final completion stage
+     - ``max_steps_reached``
+
+Per-source effective contribution was:
+
+.. list-table::
+   :header-rows: 1
+   :widths: 16 16 16 16 18 18
+
+   * - Source
+     - Scheduled
+     - Filtered
+     - Trained
+     - Trained share
+     - Token share
+   * - ``alpaca``
+     - 675
+     - 142
+     - 533
+     - 53.3%
+     - 62.32%
+   * - ``stanford``
+     - 366
+     - 5
+     - 361
+     - 36.1%
+     - 24.92%
+   * - ``gpt4``
+     - 130
+     - 24
+     - 106
+     - 10.6%
+     - 12.76%
+
+The source acceptance rates differ: Stanford Alpaca filtered fewer rows than
+the other two sources. This explains why trained-row proportions differ from
+scheduled sample weights. Token proportions differ again because accepted
+response lengths are source-dependent. The progress event makes both effects
+observable rather than silently claiming that scheduled weights equal loss
+contribution.
+
+.. figure:: ../_static/kaggle-dataset-mixing/kaggle-11-500-step-result.png
+   :alt: Final training statistics, source progress, and completion of the 500-step run
+   :width: 100%
+
+   The final 500-step log contains step 499 training statistics, source-level
+   mixing progress, and ``step=500 stage=max_steps_reached``. It is the
+   evidence for scheduled, filtered, trained, and token counts after
+   token-length filtering.
+
+10. Evaluate training stability
+-------------------------------
+
+Single-batch loss is noisy because batch size is two, so windowed statistics
+are more informative than comparing only step zero and step 499.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 34 22 22
+
+   * - Metric
+     - First 50 steps
+     - Last 50 steps
+   * - Mean loss
+     - 2.088
+     - 1.530
+   * - Median loss
+     - 2.009
+     - 1.521
+   * - Mean gradient norm
+     - 44.90
+     - 26.92
+   * - Mean train time per step
+     - 0.927 s
+     - 0.792 s
+
+The first-to-last window mean loss decreased by approximately 26.7%. All
+reported losses, gradients, learning rates, and timings remained finite. This
+supports execution stability, but it must not be presented as a downstream
+quality benchmark because no held-out evaluation set was used.
+
+The following notebook cell generated the compact final validation. It selects
+the event file with the most recorded loss points, associates it with the
+matching process-specific run state, configuration, and mix plan, and then
+checks the 500-step completion contract:
+
+.. code-block:: python
+
+   from pathlib import Path
+   from statistics import mean
+   from tensorboard.backend.event_processing.event_accumulator import (
+       EventAccumulator,
+   )
+   import json
+   import math
+
+   log_dir = Path("/tmp/areno/tfevent")
+   runs = []
+
+   for event_file in log_dir.rglob("events.out.tfevents.*"):
+       try:
+           accumulator = EventAccumulator(
+               str(event_file),
+               size_guidance={"scalars": 0},
+           )
+           accumulator.Reload()
+           tags = accumulator.Tags().get("scalars", [])
+           if "train/loss" in tags:
+               runs.append(
+                   (
+                       len(accumulator.Scalars("train/loss")),
+                       event_file,
+                       accumulator,
+                   )
+               )
+       except Exception:
+           pass
+
+   if not runs:
+       raise RuntimeError(f"No training events found under {log_dir}")
+
+   steps_recorded, event_file, accumulator = max(
+       runs,
+       key=lambda item: item[0],
+   )
+
+   pid = None
+   for state_path in log_dir.glob("dashboard_state.*.json"):
+       candidate_pid = state_path.stem.split(".")[-1]
+       if f".{candidate_pid}." in event_file.name:
+           pid = candidate_pid
+           break
+
+   if pid is None:
+       state_path = max(
+           log_dir.glob("dashboard_state.*.json"),
+           key=lambda path: path.stat().st_mtime,
+       )
+       pid = state_path.stem.split(".")[-1]
+   else:
+       state_path = log_dir / f"dashboard_state.{pid}.json"
+
+   config_path = log_dir / f"areno_run_config.{pid}.json"
+   plan_path = log_dir / f"dataset_mix_plan.{pid}.epoch-0.json"
+
+   with state_path.open(encoding="utf-8") as handle:
+       state = json.load(handle)
+   with config_path.open(encoding="utf-8") as handle:
+       config = json.load(handle)
+   with plan_path.open(encoding="utf-8") as handle:
+       plan = json.load(handle)
+
+   settings = {
+       item["key"]: item["value"]
+       for section in config["settings"]["sections"]
+       for item in section["items"]
+   }
+
+   def scalar_values(tag):
+       return [
+           float(event.value)
+           for event in accumulator.Scalars(tag)
+       ]
+
+   tags = accumulator.Tags()["scalars"]
+   loss = scalar_values("train/loss")
+   grad_norm = scalar_values("train/grad_norm")
+   target_tokens_mean = scalar_values("train/sft_target_tokens")
+   time_tag = (
+       "train/step_train_time_s"
+       if "train/step_train_time_s" in tags
+       else "time/train"
+   )
+   train_time = scalar_values(time_tag)
+
+   batch_size = int(settings["batch_size"])
+   max_steps = int(settings["max_steps"])
+   first_loss = mean(loss[:50])
+   last_loss = mean(loss[-50:])
+   loss_change = (last_loss / first_loss - 1) * 100
+
+   print("=== AReno third-run validation ===")
+   print("event_file:", event_file.name)
+   print("run_stage:", state.get("stage"))
+   print("state_step:", state.get("step"))
+   print("configured_max_steps:", max_steps)
+   print("steps_recorded:", steps_recorded)
+   print("batch_size:", batch_size)
+   print("trained_rows:", steps_recorded * batch_size)
+   print(
+       "target_tokens_trained:",
+       round(sum(target_tokens_mean) * batch_size),
+   )
+
+   print("\n=== Training metrics ===")
+   print("first_50_loss_mean:", round(first_loss, 6))
+   print("last_50_loss_mean:", round(last_loss, 6))
+   print("loss_mean_change:", f"{loss_change:.2f}%")
+   print("final_step_loss:", round(loss[-1], 6))
+   print(
+       "first_50_grad_norm_mean:",
+       round(mean(grad_norm[:50]), 6),
+   )
+   print(
+       "last_50_grad_norm_mean:",
+       round(mean(grad_norm[-50:]), 6),
+   )
+   print(
+       "overall_train_time_mean_s:",
+       round(mean(train_time), 6),
+   )
+   print(
+       "last_50_train_time_mean_s:",
+       round(mean(train_time[-50:]), 6),
+   )
+
+   print("\n=== Dataset mixing plan ===")
+   print("schedule_hash:", plan["schedule_hash"])
+   print("planned_rows:", plan["planned_rows"])
+   for source in plan["sources"]:
+       print(
+           source["name"],
+           "requested =", round(source["weight_requested"], 3),
+           "planned =", round(source["observed_proportion"], 3),
+           "selected =", source["rows_selected"],
+       )
+
+   passed = (
+       state.get("stage") == "max_steps_reached"
+       and state.get("step") == max_steps
+       and steps_recorded == max_steps
+       and len(target_tokens_mean) == max_steps
+       and all(math.isfinite(value) for value in loss)
+   )
+   print("\nVALIDATION:", "PASS" if passed else "FAIL")
+
+.. figure:: ../_static/kaggle-dataset-mixing/kaggle-12-500-step-summary.png
+   :alt: Programmatic validation summary for the 500-step dataset-mixing run
+   :width: 100%
+
+   Programmatic inspection of the TensorBoard event file and matching
+   run-state artifacts. It confirms all 500 recorded steps, 1,000 trained
+   rows, 77,533 target tokens, a 26.71% first-to-last-window mean loss
+   decrease, the expected schedule hash, finite metrics, and
+   ``VALIDATION: PASS``. The 60/30/10 values shown here are pre-filter plan
+   proportions; the effective post-filter shares are reported in the previous
+   figure and table.
+
+Evidence checklist
+------------------
+
+The retained submission contains these twelve screenshots:
+
+#. ``kaggle-01-source-revision.png`` — branch checkout and implementation SHA;
+#. ``kaggle-02-environment.png`` — PyTorch, CUDA, GPU, and extension import;
+#. ``kaggle-03-model-preparation.png`` — FP16 conversion cell;
+#. ``kaggle-04-fp16-checkpoint.png`` — completed model write and dtype;
+#. ``kaggle-05-token-limit-128-config.png`` — first-run command and config;
+#. ``kaggle-06-token-limit-128-result.png`` — first-run final progress;
+#. ``kaggle-07-structured-plan.png`` — schedule hash and planned source counts;
+#. ``kaggle-08-token-limit-256-config.png`` — second-run command and config;
+#. ``kaggle-09-token-limit-256-result.png`` — second-run final progress;
+#. ``kaggle-10-500-step-config.png`` — third-run command and config;
+#. ``kaggle-11-500-step-result.png`` — third-run source progress and completion;
+#. ``kaggle-12-500-step-summary.png`` — TensorBoard and artifact validation.
+
+All images are stored under ``docs/_static/kaggle-dataset-mixing/`` as
+lossless PNG files. The pair of final-run images is intentional: the raw log
+proves effective per-source contribution, while the compact summary verifies
+the complete 500-step metric series and pre-filter schedule.
+
+Artifacts and retained evidence
+-------------------------------
+
+The feature produces or references:
+
+* ``dataset_mix_plan.<pid>.epoch-<n>.json`` — sample-free per-epoch source plan
+  and hashes;
+* ``areno_run_config.<pid>.json`` — resolved run configuration;
+* ``areno_run_config.<pid>.txt`` — human-readable run configuration;
+* TensorBoard event files under the metrics directory;
+* notebook output containing plan, progress, train statistics, and completion.
+
+The validation intentionally did not retain a model checkpoint. A future
+quality-evaluation run should configure ``--save-path`` and record the saved
+checkpoint location and size separately.
+
+Known limitations
+-----------------
+
+* V1 supports map-style datasets; streaming datasets are not accepted.
+* The epoch index schedule is precomputed, so memory grows with planned rows.
+* Weights are sample-based and apply before tokenization/length filtering.
+* Source-specific filtering can change effective trained-row proportions.
+* Source-specific response lengths can change target-token/loss contribution.
+* The schedule hash covers sampler/index ordering only, not immutable
+  verification of remote dataset contents.
+* Exact deterministic replay starts at an epoch boundary. Mid-epoch optimizer
+  state and dataset-cursor resume are outside the current contract.
+* The Kaggle run validates two-GPU tensor-parallel SFT. It does not validate
+  multi-node execution, data parallelism, or model-quality generalization.
diff --git a/docs/getting-started/mlx.rst b/docs/getting-started/mlx.rst
index 2bd699f7..fdbcfe14 100644
--- a/docs/getting-started/mlx.rst
+++ b/docs/getting-started/mlx.rst
@@ -92,8 +92,12 @@ options have the largest effect on MLX unified-memory use:
    instead of retaining it for the next rollout.
 
 ``--adam-8bit``
-   Stores Adam moment state in the MLX backend's 8-bit representation. This
-   reduces optimizer memory; validate convergence for the target task.
+   Stores non-embedding Adam moments in the same block-wise dynamic 8-bit
+   representation used by the CUDA backend. Token-embedding optimizer moments
+   stay FP32, selected by parameter identity rather than name matching; model
+   weights, gradients, and forward behavior are unchanged. This reduces
+   optimizer memory for the remaining parameters; validate convergence for the
+   target task.
 
 ``--activation-checkpointing``
    Recomputes supported decoder activations during backward. It is enabled by
diff --git a/docs/index.rst b/docs/index.rst
index ae008795..b2d6e18f 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -37,6 +37,7 @@ AReno documentation
    :caption: Cookbook
 
    cookbook/math-rlvr
+   Kaggle Dataset-Mixing Validation <cookbook/kaggle-dataset-mixing-validation>
    cookbook/tictactoe-agentic-rl
    cookbook/duelgrid-visual-agent
 
diff --git a/examples/sft/mixed/README.md b/examples/sft/mixed/README.md
new file mode 100644
index 00000000..f6c7a397
--- /dev/null
+++ b/examples/sft/mixed/README.md
@@ -0,0 +1,49 @@
+# Deterministic weighted SFT dataset mix
+
+This local example mixes two already-normalized SFT sources without network
+services or external databases:
+
+```bash
+areno train \
+  --algo sft \
+  --ckpt /path/to/local/model \
+  --dataset-mix-config examples/sft/mixed/mix.json \
+  --dataset-loader-fn examples/sft/mixed/dataset_loader.py \
+  --world-size 1 \
+  --tp-size 1 \
+  --batch-size 2 \
+  --epochs 1 \
+  --metrics-log-dir outputs/mixed-sft/metrics
+```
+
+For the default seed ``42``, bounded ``cycle`` policy, and source shuffling,
+the same sources can be supplied directly on the command line:
+
+```bash
+areno train \
+  --algo sft \
+  --ckpt /path/to/local/model \
+  --dataset-source math=examples/sft/mixed/math.jsonl:0.7 \
+  --dataset-source code=examples/sft/mixed/code.jsonl:0.3 \
+  --dataset-mix-samples-per-epoch 20 \
+  --dataset-loader-fn examples/sft/mixed/dataset_loader.py \
+  --world-size 1 \
+  --tp-size 1 \
+  --batch-size 2 \
+  --epochs 1 \
+  --metrics-log-dir outputs/mixed-sft/metrics
+```
+
+Repeat ``--dataset-source NAME=PATH:WEIGHT`` at least twice. Use the JSON
+manifest form to keep the complete mix contract in a file. Additional
+``--dataset-source`` options work the same way for three or more datasets.
+
+Before model or worker initialization, AReno validates both sources and prints
+the planned source counts. It also writes
+`dataset_mix_plan.<pid>.json` under the metrics directory without including
+prompt or response contents. During training,
+`stage=dataset_mix_progress` logs scheduled, filtered, and trained rows plus
+trained target-token counts and observed row/token proportions.
+
+For a boundary-input example, change a weight to `0`; validation then fails
+with `weight must be finite and positive`.
diff --git a/examples/sft/mixed/code.jsonl b/examples/sft/mixed/code.jsonl
new file mode 100644
index 00000000..93f28103
--- /dev/null
+++ b/examples/sft/mixed/code.jsonl
@@ -0,0 +1,2 @@
+{"prompt":"Write a Python expression that adds one to x.","response":"x + 1"}
+{"prompt":"Write a Python expression that returns the length of items.","response":"len(items)"}
diff --git a/examples/sft/mixed/dataset_loader.py b/examples/sft/mixed/dataset_loader.py
new file mode 100644
index 00000000..7573234d
--- /dev/null
+++ b/examples/sft/mixed/dataset_loader.py
@@ -0,0 +1,5 @@
+"""Pass through already-normalized rows for the deterministic mix example."""
+
+
+def load_training_dataset(dataset_path, *, default_loader, **_kwargs):
+    return default_loader(dataset_path)
diff --git a/examples/sft/mixed/math.jsonl b/examples/sft/mixed/math.jsonl
new file mode 100644
index 00000000..7e57d6f3
--- /dev/null
+++ b/examples/sft/mixed/math.jsonl
@@ -0,0 +1,3 @@
+{"prompt":"What is 2 + 2?","response":"4"}
+{"prompt":"What is 3 × 3?","response":"9"}
+{"prompt":"What is 12 ÷ 4?","response":"3"}
diff --git a/examples/sft/mixed/mix.json b/examples/sft/mixed/mix.json
new file mode 100644
index 00000000..a8d1b40d
--- /dev/null
+++ b/examples/sft/mixed/mix.json
@@ -0,0 +1,19 @@
+{
+  "version": 1,
+  "seed": 42,
+  "exhaustion": "cycle",
+  "shuffle_within_sources": true,
+  "samples_per_epoch": 20,
+  "sources": [
+    {
+      "name": "math",
+      "path": "math.jsonl",
+      "weight": 0.7
+    },
+    {
+      "name": "code",
+      "path": "code.jsonl",
+      "weight": 0.3
+    }
+  ]
+}
diff --git a/tests/test_adamw_8bit_blockwise_cpu.py b/tests/test_adamw_8bit_blockwise_cpu.py
index c1a4fd9d..902d56a1 100644
--- a/tests/test_adamw_8bit_blockwise_cpu.py
+++ b/tests/test_adamw_8bit_blockwise_cpu.py
@@ -1,9 +1,24 @@
 from __future__ import annotations
 
+import copy
+from types import SimpleNamespace
+from unittest.mock import patch
+
 import torch
 
-from areno.engine.optim import AdamW8bit
-from areno.engine.optim.adamw_8bit import _dequantize_positive, _quantize_positive
+from areno.engine.optim import AdamW8bit, set_optimizer_state_precision
+from areno.engine.optim.adamw_8bit import (
+    _dequantize_positive,
+    _dequantize_symmetric,
+    _quantize_positive,
+    _quantize_symmetric,
+)
+from areno.engine.optim.dynamic_quant import (
+    SIGNED_DYNAMIC_MAP,
+    SIGNED_DYNAMIC_ZERO,
+    UNSIGNED_DYNAMIC_MAP,
+    UNSIGNED_DYNAMIC_ZERO,
+)
 
 
 def test_adamw8bit_uses_parameter_local_block_scales() -> None:
@@ -30,13 +45,37 @@ def test_adamw8bit_uses_parameter_local_block_scales() -> None:
     torch.testing.assert_close(second, torch.tensor([-1.0e-3, 1.0e-3, -1.0e-3]), atol=1.0e-6, rtol=0.0)
 
 
-def test_adamw8bit_second_moment_uses_full_linear_range_per_block() -> None:
-    values = torch.tensor([0.0, 1.0 / 255.0, 128.0 / 255.0, 1.0])
+def test_adamw8bit_dynamic_codebooks_match_paper_reference_construction() -> None:
+    assert len(SIGNED_DYNAMIC_MAP) == len(UNSIGNED_DYNAMIC_MAP) == 256
+    assert SIGNED_DYNAMIC_ZERO == 127
+    assert UNSIGNED_DYNAMIC_ZERO == 0
+    assert SIGNED_DYNAMIC_MAP[0] == -0.99296875
+    assert SIGNED_DYNAMIC_MAP[-1] == UNSIGNED_DYNAMIC_MAP[-1] == 1.0
+    assert all(left <= right for left, right in zip(SIGNED_DYNAMIC_MAP, SIGNED_DYNAMIC_MAP[1:]))
+    assert all(left <= right for left, right in zip(UNSIGNED_DYNAMIC_MAP, UNSIGNED_DYNAMIC_MAP[1:]))
+
+
+def test_adamw8bit_dynamic_codebook_golden_round_trip() -> None:
+    signed = torch.tensor(SIGNED_DYNAMIC_MAP)
+    unsigned = torch.tensor(UNSIGNED_DYNAMIC_MAP)
+
+    signed_q, signed_scale = _quantize_symmetric(signed)
+    unsigned_q, unsigned_scale = _quantize_positive(unsigned)
+
+    expected_codes = torch.arange(256, dtype=torch.int64).to(torch.uint8)
+    assert torch.equal(signed_q, expected_codes)
+    assert torch.equal(unsigned_q, expected_codes)
+    torch.testing.assert_close(_dequantize_symmetric(signed_q, signed_scale), signed)
+    torch.testing.assert_close(_dequantize_positive(unsigned_q, unsigned_scale), unsigned)
+
+
+def test_adamw8bit_unsigned_dynamic_map_preserves_small_second_moments() -> None:
+    values = torch.tensor([0.0, 7.75e-7, 5.5e-5, 5.5e-3, 0.55, 1.0])
 
     quantized, scale = _quantize_positive(values)
     restored = _dequantize_positive(quantized, scale)
 
-    torch.testing.assert_close(restored, values)
+    torch.testing.assert_close(restored, values, rtol=0.15, atol=1.0e-8)
 
 
 def test_adamw8bit_same_lr_does_not_amplify_constant_gradient_step() -> None:
@@ -53,3 +92,229 @@ def test_adamw8bit_same_lr_does_not_amplify_constant_gradient_step() -> None:
     optimizer.step()
 
     torch.testing.assert_close(parameter, torch.full_like(parameter, 0.999), rtol=0.0, atol=1.0e-6)
+
+
+def test_adamw8bit_routes_embedding_role_to_fp32_state_without_name_matching() -> None:
+    embedding = torch.nn.Parameter(torch.ones(12))
+    embedding._areno_optimizer_role = "token_embedding"
+    ordinary = torch.nn.Parameter(torch.ones(12))
+    optimizer = AdamW8bit(
+        [embedding, ordinary],
+        lr=1.0e-3,
+        betas=(0.9, 0.99),
+        weight_decay=0.0,
+        bucket_numel=64,
+        quant_block_size=4,
+    )
+
+    assert [state.precision for state in optimizer._states] == ["fp32", "8bit"]
+    embedding.grad = torch.linspace(-1.0, 1.0, 12)
+    ordinary.grad = embedding.grad.clone()
+    optimizer.step()
+    state = optimizer.state_dict()
+
+    assert state["quantizer"] == "dynamic-tree-v1"
+    assert [item["precision"] for item in state["precision_policy"]] == ["fp32", "8bit"]
+    assert state["precision_policy"][0]["role"] == "token_embedding"
+    assert state["state"][0]["exp_avg"] is not None
+    assert state["state"][0]["exp_avg_q"] is None
+    assert state["state"][1]["exp_avg"] is None
+    assert state["state"][1]["exp_avg_q"].dtype == torch.uint8
+
+
+def test_vocab_parallel_embedding_role_is_safe_for_pretrained_weights_and_dp_sharding() -> None:
+    from areno.engine.layers.vocab import VocabParallelEmbedding
+
+    with patch(
+        "areno.engine.layers.vocab.get_tp_context",
+        return_value=SimpleNamespace(rank=1, world_size=2),
+    ):
+        embedding = VocabParallelEmbedding(18, 6)
+    loaded_weight = torch.linspace(-1.0, 1.0, embedding.weight.numel()).reshape_as(embedding.weight)
+    embedding.weight.data.copy_(loaded_weight)
+    optimizer = AdamW8bit(
+        embedding.parameters(),
+        lr=1.0e-3,
+        betas=(0.9, 0.99),
+        weight_decay=0.0,
+        dp_rank=1,
+        dp_size=2,
+    )
+
+    assert embedding.weight._areno_optimizer_role == "token_embedding"
+    assert optimizer._states[0].precision == "fp32"
+    assert optimizer.buckets[0].shard_numel == (embedding.weight.numel() + 1) // 2
+    torch.testing.assert_close(embedding.weight, loaded_weight, rtol=0.0, atol=0.0)
+
+
+def test_adamw8bit_explicit_precision_override_beats_embedding_default_and_deduplicates_ties() -> None:
+    tied = torch.nn.Parameter(torch.ones(8))
+    tied._areno_optimizer_role = "token_embedding"
+    set_optimizer_state_precision(tied, "8bit")
+    fp32 = torch.nn.Parameter(torch.ones(4))
+    optimizer = AdamW8bit(
+        [
+            {"params": [tied], "state_precision": "fp32"},
+            {"params": [tied, fp32], "state_precision": "fp32"},
+        ],
+        lr=1.0e-3,
+        betas=(0.9, 0.99),
+        weight_decay=0.0,
+        bucket_numel=64,
+    )
+
+    assert optimizer.model_params == [tied, fp32]
+    assert [state.precision for state in optimizer._states] == ["8bit", "fp32"]
+
+
+def test_adamw8bit_mixed_state_checkpoint_round_trip_preserves_next_update() -> None:
+    embedding = torch.nn.Parameter(torch.linspace(-0.5, 0.5, 12))
+    embedding._areno_optimizer_role = "token_embedding"
+    ordinary = torch.nn.Parameter(torch.linspace(0.25, -0.25, 9))
+    first = AdamW8bit(
+        [embedding, ordinary],
+        lr=4.0e-4,
+        betas=(0.9, 0.99),
+        weight_decay=0.01,
+        bucket_numel=64,
+        quant_block_size=4,
+    )
+    embedding.grad = torch.linspace(-0.7, 0.3, embedding.numel())
+    ordinary.grad = torch.linspace(0.4, -0.8, ordinary.numel())
+    first.step()
+    checkpoint = copy.deepcopy(first.state_dict())
+
+    restored_embedding = torch.nn.Parameter(embedding.detach().clone())
+    restored_embedding._areno_optimizer_role = "token_embedding"
+    restored_ordinary = torch.nn.Parameter(ordinary.detach().clone())
+    restored = AdamW8bit(
+        [restored_embedding, restored_ordinary],
+        lr=4.0e-4,
+        betas=(0.9, 0.99),
+        weight_decay=0.01,
+        bucket_numel=64,
+        quant_block_size=4,
+    )
+    restored.load_state_dict(checkpoint)
+
+    next_embedding_grad = torch.linspace(0.6, -0.2, embedding.numel())
+    next_ordinary_grad = torch.linspace(-0.1, 0.9, ordinary.numel())
+    embedding.grad = next_embedding_grad.clone()
+    restored_embedding.grad = next_embedding_grad.clone()
+    ordinary.grad = next_ordinary_grad.clone()
+    restored_ordinary.grad = next_ordinary_grad.clone()
+    first.step()
+    restored.step()
+
+    torch.testing.assert_close(restored_embedding, embedding, rtol=0.0, atol=0.0)
+    torch.testing.assert_close(restored_ordinary, ordinary, rtol=0.0, atol=0.0)
+    for actual, expected in zip(restored.state_dict()["state"], first.state_dict()["state"], strict=True):
+        for key in ("exp_avg_q", "exp_avg_scale", "exp_avg_sq_q", "exp_avg_sq_scale", "exp_avg", "exp_avg_sq"):
+            if expected[key] is None:
+                assert actual[key] is None
+            else:
+                torch.testing.assert_close(actual[key], expected[key], rtol=0.0, atol=0.0)
+
+
+def test_adamw8bit_checkpoint_restores_saved_precision_policy_by_parameter_identity() -> None:
+    embedding = torch.nn.Parameter(torch.ones(6))
+    embedding._areno_optimizer_role = "token_embedding"
+    ordinary = torch.nn.Parameter(torch.ones(6))
+    source = AdamW8bit(
+        [embedding, ordinary],
+        lr=1.0e-3,
+        betas=(0.9, 0.99),
+        weight_decay=0.0,
+        bucket_numel=16,
+        quant_block_size=4,
+    )
+    embedding.grad = torch.ones_like(embedding)
+    ordinary.grad = torch.ones_like(ordinary)
+    source.step()
+
+    restored_embedding = torch.nn.Parameter(embedding.detach().clone())
+    restored_ordinary = torch.nn.Parameter(ordinary.detach().clone())
+    restored = AdamW8bit(
+        [restored_embedding, restored_ordinary],
+        lr=1.0e-3,
+        betas=(0.9, 0.99),
+        weight_decay=0.0,
+        bucket_numel=16,
+        quant_block_size=4,
+    )
+    assert [state.precision for state in restored._states] == ["8bit"]
+
+    restored.load_state_dict(source.state_dict())
+
+    assert [state.precision for state in restored._states] == ["fp32", "8bit"]
+    assert restored._parameter_roles[id(restored_embedding)] == "token_embedding"
+
+
+def test_adamw8bit_reports_mixed_state_storage() -> None:
+    embedding = torch.nn.Parameter(torch.ones(12))
+    embedding._areno_optimizer_role = "token_embedding"
+    ordinary = torch.nn.Parameter(torch.ones(12))
+    optimizer = AdamW8bit(
+        [embedding, ordinary],
+        lr=1.0e-3,
+        betas=(0.9, 0.99),
+        weight_decay=0.0,
+        bucket_numel=64,
+        quant_block_size=4,
+    )
+    embedding.grad = torch.ones_like(embedding)
+    ordinary.grad = torch.ones_like(ordinary)
+    optimizer.step()
+
+    assert optimizer.state_memory_metrics() == {
+        "quantized_state_bytes": 24,
+        "fp32_exempt_bytes": 96,
+        "block_metadata_bytes": 24,
+        "total_bytes": 144,
+    }
+
+
+def test_adamw8bit_nonfinite_gradient_skips_only_affected_block() -> None:
+    parameter = torch.nn.Parameter(torch.zeros(8))
+    optimizer = AdamW8bit(
+        [parameter],
+        lr=1.0e-3,
+        betas=(0.9, 0.99),
+        weight_decay=0.0,
+        bucket_numel=16,
+        quant_block_size=4,
+    )
+    gradient = torch.ones_like(parameter)
+    gradient[1] = torch.inf
+    parameter.grad = gradient
+    optimizer.step()
+
+    torch.testing.assert_close(parameter[:4], torch.zeros(4))
+    assert torch.all(parameter[4:] < 0)
+
+
+def test_adamw8bit_disk_offload_supports_mixed_state(tmp_path) -> None:
+    embedding = torch.nn.Parameter(torch.ones(8))
+    embedding._areno_optimizer_role = "token_embedding"
+    ordinary = torch.nn.Parameter(torch.ones(8))
+    candidate = AdamW8bit(
+        [embedding, ordinary],
+        lr=1.0e-3,
+        betas=(0.9, 0.99),
+        weight_decay=0.0,
+        bucket_numel=16,
+        quant_block_size=4,
+    )
+    candidate.configure_state_offload(mode="disk", directory=str(tmp_path), batch_size=2)
+    embedding.grad = torch.linspace(-1.0, 1.0, embedding.numel())
+    ordinary.grad = torch.linspace(1.0, -1.0, ordinary.numel())
+    candidate.step()
+
+    assert all(state.offload_file is not None for state in candidate._states)
+    saved = candidate.state_dict()["state"]
+    assert saved[0]["exp_avg"] is not None and saved[0]["exp_avg_q"] is None
+    assert saved[1]["exp_avg"] is None and saved[1]["exp_avg_q"] is not None
+    candidate.onload_state(torch.device("cpu"))
+    assert candidate._states[0].exp_avg is not None
+    assert candidate._states[1].exp_avg_q is not None
+    assert not list(tmp_path.rglob("*.mmap"))
diff --git a/tests/test_agent_skills_cpu.py b/tests/test_agent_skills_cpu.py
index a081da37..5e796f74 100644
--- a/tests/test_agent_skills_cpu.py
+++ b/tests/test_agent_skills_cpu.py
@@ -26,7 +26,7 @@ def test_repository_agent_skills_are_valid():
         )
 
     result = json.loads(process.stdout)
-    assert result["skill_count"] == 10
+    assert result["skill_count"] == 11
     assert result["script_count"] >= 15
 
 
diff --git a/tests/test_dataset_mixing_cpu.py b/tests/test_dataset_mixing_cpu.py
new file mode 100644
index 00000000..dd6083f2
--- /dev/null
+++ b/tests/test_dataset_mixing_cpu.py
@@ -0,0 +1,616 @@
+from __future__ import annotations
+
+import json
+import math
+from pathlib import Path
+
+import pytest
+
+from areno.api.data import DATASET_MIX_METADATA_KEY, DatasetMixSource, WeightedMixedDataset
+from areno.cli import train as train_cli
+
+
+def _source(name: str, size: int, weight: float) -> DatasetMixSource:
+    return DatasetMixSource(
+        name=name,
+        dataset=[{"prompt": f"{name}-{index}", "response": str(index)} for index in range(size)],
+        weight=weight,
+    )
+
+
+def _identity(dataset: WeightedMixedDataset) -> list[tuple[str, int, int]]:
+    return [
+        (
+            row[DATASET_MIX_METADATA_KEY]["source"],
+            row[DATASET_MIX_METADATA_KEY]["source_index"],
+            row[DATASET_MIX_METADATA_KEY]["cycle"],
+        )
+        for row in dataset
+    ]
+
+
+def test_weighted_mix_is_deterministic_for_seed_and_epoch():
+    sources = [_source("math", 5, 0.7), _source("code", 4, 0.3)]
+    first = WeightedMixedDataset(sources, seed=42, exhaustion="renormalize")
+    second = WeightedMixedDataset(sources, seed=42, exhaustion="renormalize")
+
+    assert _identity(first) == _identity(second)
+    assert first.summary()["schedule_hash"] == second.summary()["schedule_hash"]
+
+    first.set_epoch(1)
+    assert _identity(first) != _identity(second)
+
+
+@pytest.mark.parametrize("epoch", [True, 1.5, "1", -1])
+def test_weighted_mix_rejects_invalid_epoch(epoch):
+    dataset = WeightedMixedDataset(
+        [_source("first", 2, 1), _source("second", 2, 1)],
+        seed=42,
+        exhaustion="renormalize",
+    )
+
+    with pytest.raises(ValueError, match="non-negative integer"):
+        dataset.set_epoch(epoch)
+
+
+def test_renormalize_emits_every_record_once():
+    dataset = WeightedMixedDataset(
+        [_source("small", 2, 0.8), _source("large", 7, 0.2)],
+        seed=7,
+        exhaustion="renormalize",
+        shuffle_within_sources=False,
+    )
+
+    identities = _identity(dataset)
+    assert len(identities) == 9
+    assert len({(source, index) for source, index, _cycle in identities}) == 9
+    assert all(cycle == 0 for _source_name, _index, cycle in identities)
+    assert dataset.summary()["termination_reason"] == "all_sources_exhausted"
+    assert all(source["duplicates"] == 0 for source in dataset.summary()["sources"])
+
+
+def test_stop_never_repeats_records_and_may_leave_rows_unused():
+    dataset = WeightedMixedDataset(
+        [_source("small", 1, 0.9), _source("large", 8, 0.1)],
+        seed=3,
+        exhaustion="stop",
+        shuffle_within_sources=False,
+    )
+
+    identities = _identity(dataset)
+    assert len(identities) < 9
+    assert len({(source, index) for source, index, _cycle in identities}) == len(identities)
+    assert dataset.summary()["termination_reason"].startswith("source_exhausted:")
+
+
+def test_cycle_visits_every_record_and_can_repeat_small_sources():
+    dataset = WeightedMixedDataset(
+        [_source("small", 1, 0.9), _source("large", 4, 0.1)],
+        seed=11,
+        exhaustion="cycle",
+        shuffle_within_sources=False,
+    )
+
+    identities = _identity(dataset)
+    unique = {(source, index) for source, index, _cycle in identities}
+    assert unique == {("small", 0), ("large", 0), ("large", 1), ("large", 2), ("large", 3)}
+    assert len(identities) > len(unique)
+    assert dataset.summary()["termination_reason"] == "all_sources_exhausted_once"
+    assert next(item for item in dataset.summary()["sources"] if item["name"] == "small")["duplicates"] > 0
+
+
+def test_samples_per_epoch_bounds_cycle():
+    dataset = WeightedMixedDataset(
+        [_source("small", 1, 0.99), _source("rare", 2, 0.01)],
+        seed=5,
+        exhaustion="cycle",
+        samples_per_epoch=3,
+    )
+
+    assert len(dataset) == 3
+    assert dataset.summary()["termination_reason"] == "samples_per_epoch"
+
+
+def test_samples_per_epoch_is_rejected_for_non_cycle_policies():
+    with pytest.raises(ValueError, match="only supported with exhaustion='cycle'"):
+        WeightedMixedDataset(
+            [_source("first", 1, 1.0), _source("second", 1, 1.0)],
+            seed=5,
+            exhaustion="renormalize",
+            samples_per_epoch=1,
+        )
+
+
+def test_observed_proportions_follow_weights_with_documented_tolerance():
+    dataset = WeightedMixedDataset(
+        [_source("major", 10_000, 0.7), _source("minor", 10_000, 0.3)],
+        seed=42,
+        exhaustion="cycle",
+        shuffle_within_sources=False,
+        samples_per_epoch=2_000,
+    )
+
+    observed = {source["name"]: source["observed_proportion"] for source in dataset.summary()["sources"]}
+    assert observed["major"] == pytest.approx(0.7, abs=0.05)
+    assert observed["minor"] == pytest.approx(0.3, abs=0.05)
+
+
+def test_many_source_cycle_uses_fixed_budget_and_reports_rare_source_risk():
+    dataset = WeightedMixedDataset(
+        [
+            _source("math", 3, 0.6),
+            _source("code", 2, 0.3),
+            _source("general", 4, 0.099),
+            _source("rare", 1, 0.001),
+        ],
+        seed=42,
+        exhaustion="cycle",
+        samples_per_epoch=100,
+    )
+
+    summary = dataset.summary()
+    assert len(dataset) == 100
+    assert summary["sampler_version"] == 1
+    assert summary["weight_unit"] == "sample"
+    assert summary["termination_reason"] == "samples_per_epoch"
+    assert summary["mix_spec_hash"].startswith("sha256:")
+    assert summary["schedule_hash"].startswith("sha256:")
+    assert any("source 'rare'" in warning for warning in summary["warnings"])
+
+
+def test_weighted_mix_summary_returns_defensive_nested_copies():
+    dataset = WeightedMixedDataset(
+        [_source("common", 2, 0.999), _source("rare", 1, 0.001)],
+        seed=42,
+        exhaustion="cycle",
+        samples_per_epoch=100,
+    )
+
+    summary = dataset.summary()
+    summary["warnings"].append("caller mutation")
+    summary["sources"][0]["name"] = "caller mutation"
+
+    fresh_summary = dataset.summary()
+    assert "caller mutation" not in fresh_summary["warnings"]
+    assert fresh_summary["sources"][0]["name"] == "common"
+
+
+@pytest.mark.parametrize("weight", [0, -1, math.nan, math.inf])
+def test_weighted_mix_rejects_invalid_weights(weight):
+    with pytest.raises(ValueError, match="weight must be finite and positive"):
+        WeightedMixedDataset(
+            [_source("bad", 1, weight), _source("good", 1, 1.0)],
+            seed=1,
+            exhaustion="stop",
+        )
+
+
+def test_weighted_mix_rejects_reserved_metadata_field():
+    with pytest.raises(ValueError, match="reserved field"):
+        WeightedMixedDataset(
+            [
+                DatasetMixSource("bad", [{DATASET_MIX_METADATA_KEY: {}, "prompt": "x"}], 1.0),
+                _source("good", 1, 1.0),
+            ],
+            seed=1,
+            exhaustion="stop",
+        )
+
+
+def test_weighted_mix_rejects_bool_weight_and_non_bool_shuffle():
+    with pytest.raises(ValueError, match="weight must be finite and positive"):
+        WeightedMixedDataset(
+            [_source("good", 1, 1.0), DatasetMixSource("bad", [{"prompt": "x"}], True)],
+            seed=1,
+            exhaustion="stop",
+        )
+    with pytest.raises(ValueError, match="shuffle_within_sources must be a boolean"):
+        WeightedMixedDataset(
+            [_source("first", 1, 1.0), _source("second", 1, 1.0)],
+            seed=1,
+            exhaustion="stop",
+            shuffle_within_sources=1,
+        )
+
+
+def test_weighted_mix_normalizes_large_finite_weights_without_overflow():
+    dataset = WeightedMixedDataset(
+        [_source("first", 2, 1e308), _source("second", 2, 1e308)],
+        seed=1,
+        exhaustion="renormalize",
+    )
+
+    assert [source["weight_requested"] for source in dataset.summary()["sources"]] == [0.5, 0.5]
+
+
+def test_weighted_mix_rejects_unrepresentable_weight_range():
+    with pytest.raises(ValueError, match="unsupported numeric range"):
+        WeightedMixedDataset(
+            [_source("tiny", 1, 1e-300), _source("huge", 1, 1e300)],
+            seed=1,
+            exhaustion="renormalize",
+        )
+
+
+def test_weighted_mix_rejects_weights_that_underflow_during_final_normalization():
+    with pytest.raises(ValueError, match="unsupported numeric range"):
+        WeightedMixedDataset(
+            [
+                _source("smallest", 1, 5e-324),
+                _source("large-a", 1, 1.0),
+                _source("large-b", 1, 1.0),
+            ],
+            seed=1,
+            exhaustion="renormalize",
+        )
+
+
+def test_weighted_mix_rejects_non_printable_source_names():
+    with pytest.raises(ValueError, match="non-printable"):
+        WeightedMixedDataset(
+            [_source("safe", 1, 1), _source("forged\nlog", 1, 1)],
+            seed=1,
+            exhaustion="renormalize",
+        )
+
+
+def test_weighted_mix_rejects_ambiguous_source_name_whitespace():
+    with pytest.raises(ValueError, match="surrounding whitespace"):
+        WeightedMixedDataset(
+            [_source("safe", 1, 1), _source(" padded ", 1, 1)],
+            seed=1,
+            exhaustion="renormalize",
+        )
+
+
+def test_mix_manifest_loads_two_local_sft_sources_through_shared_loader(tmp_path):
+    first = tmp_path / "first.jsonl"
+    second = tmp_path / "second.jsonl"
+    first.write_text('{"instruction":"math","output":"1"}\n', encoding="utf-8")
+    second.write_text('{"instruction":"code","output":"2"}\n', encoding="utf-8")
+    loader = tmp_path / "loader.py"
+    loader.write_text(
+        "def load_training_dataset(dataset_path, *, default_loader, **kwargs):\n"
+        "    return [\n"
+        "        {'prompt': row['instruction'], 'response': row['output']}\n"
+        "        for row in default_loader(dataset_path)\n"
+        "    ]\n",
+        encoding="utf-8",
+    )
+    manifest = tmp_path / "mix.json"
+    manifest.write_text(
+        json.dumps(
+            {
+                "version": 1,
+                "seed": 42,
+                "exhaustion": "renormalize",
+                "sources": [
+                    {"name": "math", "path": first.name, "weight": 7},
+                    {"name": "code", "path": second.name, "weight": 3},
+                ],
+            }
+        ),
+        encoding="utf-8",
+    )
+
+    def load_dataset(_builder, *, data_files, split):
+        assert split == "train"
+        return [json.loads(line) for line in Path(data_files).read_text(encoding="utf-8").splitlines()]
+
+    dataset = train_cli._load_mixed_dataset_for_training(
+        str(manifest),
+        model_hub="hf",
+        dataset_loader_fn=str(loader),
+        load_dataset=load_dataset,
+        load_from_disk=lambda _path: None,
+    )
+
+    assert len(dataset) == 2
+    assert {row["prompt"] for row in dataset} == {"math", "code"}
+    assert {row[DATASET_MIX_METADATA_KEY]["source"] for row in dataset} == {"math", "code"}
+
+
+def test_dataset_source_shorthand_preserves_remote_ref_colons():
+    manifest = train_cli._dataset_mix_manifest_from_sources(
+        (
+            "math=gsm8k:main:train:0.7",
+            "code=org/code-dataset:default:train:0.3",
+        )
+    )
+
+    assert manifest["seed"] == 42
+    assert manifest["exhaustion"] == "cycle"
+    assert manifest["shuffle_within_sources"] is True
+    assert manifest["samples_per_epoch"] is None
+    assert manifest["sources"] == [
+        {"name": "math", "path": "gsm8k:main:train", "weight": 0.7},
+        {"name": "code", "path": "org/code-dataset:default:train", "weight": 0.3},
+    ]
+
+
+def test_mix_source_path_resolution_distinguishes_local_and_remote_refs(tmp_path):
+    config_path = tmp_path / "mix.json"
+    existing = tmp_path / "data" / "train.jsonl"
+    existing.parent.mkdir()
+    existing.write_text("", encoding="utf-8")
+
+    assert train_cli._resolve_mix_source_path(config_path, "data/train.jsonl") == str(existing)
+    assert train_cli._resolve_mix_source_path(config_path, "./missing.jsonl") == str(tmp_path / "missing.jsonl")
+    assert train_cli._resolve_mix_source_path(config_path, "org/remote-dataset") == "org/remote-dataset"
+    assert train_cli._resolve_mix_source_path(config_path, "org/remote:config:train") == "org/remote:config:train"
+
+
+@pytest.mark.parametrize(
+    ("source_specs", "message"),
+    [
+        (("one=dataset:1",), "at least twice"),
+        (("broken", "two=dataset:1"), "NAME=PATH:WEIGHT"),
+        (("same=one:1", "same=two:1"), "duplicate source name"),
+        (("one=dataset:nan", "two=dataset:1"), "finite and positive"),
+    ],
+)
+def test_dataset_source_shorthand_rejects_invalid_entries(source_specs, message):
+    with pytest.raises(ValueError, match=message):
+        train_cli._dataset_mix_manifest_from_sources(source_specs)
+
+
+def test_dataset_source_shorthand_rejects_log_injection_name():
+    with pytest.raises(ValueError, match="non-printable"):
+        train_cli._dataset_mix_manifest_from_sources(("safe=org/first:1", "forged\nsummary=org/second:1"))
+
+
+def test_dataset_source_shorthand_loads_remote_sources_through_shared_loader(tmp_path):
+    loader = tmp_path / "loader.py"
+    loader.write_text(
+        "def load_training_dataset(dataset_path, *, default_loader, **kwargs):\n"
+        "    return [\n"
+        "        {'prompt': row['instruction'], 'response': row['output']}\n"
+        "        for row in default_loader(dataset_path)\n"
+        "    ]\n",
+        encoding="utf-8",
+    )
+    loaded = []
+
+    def load_dataset(name, **_kwargs):
+        loaded.append(name)
+        return [{"instruction": name, "output": "ok"}]
+
+    dataset = train_cli._load_dataset_sources_for_training(
+        ("first=org/first:0.7", "second=org/second:0.3"),
+        samples_per_epoch=100,
+        model_hub="hf",
+        dataset_loader_fn=str(loader),
+        load_dataset=load_dataset,
+        load_from_disk=lambda _path: None,
+    )
+
+    assert loaded == ["org/first", "org/second"]
+    assert len(dataset) == 100
+    assert {row["prompt"] for row in dataset} == {"org/first", "org/second"}
+    assert {row[DATASET_MIX_METADATA_KEY]["source"] for row in dataset} == {"first", "second"}
+
+
+@pytest.mark.parametrize("version", [True, 1.0, 2])
+def test_mix_manifest_requires_integer_version_one(tmp_path, version):
+    manifest = tmp_path / "mix.json"
+    manifest.write_text(json.dumps({"version": version}), encoding="utf-8")
+
+    with pytest.raises(ValueError, match="version must be 1"):
+        train_cli._read_dataset_mix_manifest(manifest)
+
+
+@pytest.mark.parametrize(
+    "payload",
+    [
+        {
+            "version": 1,
+            "seed": 1,
+            "exhaustion": "renormalize",
+            "shufle_within_sources": False,
+            "sources": [
+                {"name": "first", "path": "first", "weight": 1},
+                {"name": "second", "path": "second", "weight": 1},
+            ],
+        },
+        {
+            "version": 1,
+            "seed": 1,
+            "exhaustion": "renormalize",
+            "sources": [
+                {"name": "first", "path": "first", "weight": 1, "weigth": 2},
+                {"name": "second", "path": "second", "weight": 1},
+            ],
+        },
+    ],
+)
+def test_mix_manifest_rejects_unknown_fields_instead_of_ignoring_typos(tmp_path, payload):
+    manifest = tmp_path / "mix.json"
+    manifest.write_text(json.dumps(payload), encoding="utf-8")
+
+    with pytest.raises(ValueError, match="unsupported field"):
+        train_cli._read_dataset_mix_manifest(manifest)
+
+
+def test_mix_manifest_normalizes_source_name_and_path_whitespace(tmp_path):
+    manifest = tmp_path / "mix.json"
+    manifest.write_text(
+        json.dumps(
+            {
+                "version": 1,
+                "seed": 1,
+                "exhaustion": "renormalize",
+                "sources": [
+                    {"name": " first ", "path": " org/first ", "weight": 1},
+                    {"name": "second", "path": "org/second", "weight": 1},
+                ],
+            }
+        ),
+        encoding="utf-8",
+    )
+
+    parsed = train_cli._read_dataset_mix_manifest(manifest)
+
+    assert parsed["sources"][0] == {"name": "first", "path": "org/first", "weight": 1.0}
+
+
+def test_mix_manifest_rejects_final_weight_normalization_underflow(tmp_path):
+    manifest = tmp_path / "mix.json"
+    manifest.write_text(
+        json.dumps(
+            {
+                "version": 1,
+                "seed": 1,
+                "exhaustion": "renormalize",
+                "sources": [
+                    {"name": "smallest", "path": "smallest", "weight": 5e-324},
+                    {"name": "large-a", "path": "large-a", "weight": 1.0},
+                    {"name": "large-b", "path": "large-b", "weight": 1.0},
+                ],
+            }
+        ),
+        encoding="utf-8",
+    )
+
+    with pytest.raises(ValueError, match="unsupported numeric range"):
+        train_cli._read_dataset_mix_manifest(manifest)
+
+
+def test_mix_loader_reports_incompatible_source_without_sample_contents(tmp_path):
+    manifest = tmp_path / "mix.json"
+    manifest.write_text(
+        json.dumps(
+            {
+                "version": 1,
+                "seed": 1,
+                "exhaustion": "stop",
+                "sources": [
+                    {"name": "good", "path": "good", "weight": 1},
+                    {"name": "bad", "path": "bad", "weight": 1},
+                ],
+            }
+        ),
+        encoding="utf-8",
+    )
+
+    def load_dataset(name, **_kwargs):
+        if name == "good":
+            return [{"prompt": "safe", "response": "ok"}]
+        return [{"prompt": "do-not-log-this-secret"}]
+
+    with pytest.raises(ValueError, match=r"stage=dataset_mix_validation source=bad.*missing required SFT field"):
+        train_cli._load_mixed_dataset_for_training(
+            str(manifest),
+            model_hub="hf",
+            dataset_loader_fn=None,
+            load_dataset=load_dataset,
+            load_from_disk=lambda _path: None,
+        )
+
+
+def test_dataset_source_shorthand_rejects_log_injection_path():
+    with pytest.raises(ValueError, match="non-printable"):
+        train_cli._dataset_mix_manifest_from_sources(("safe=org/first:1", "second=org/forged\npath:1"))
+
+
+def test_mix_validation_checks_later_rows_before_backend_initialization(tmp_path):
+    manifest = tmp_path / "mix.json"
+    manifest.write_text(
+        json.dumps(
+            {
+                "version": 1,
+                "seed": 1,
+                "exhaustion": "renormalize",
+                "sources": [
+                    {"name": "good", "path": "good", "weight": 1},
+                    {"name": "bad", "path": "bad", "weight": 1},
+                ],
+            }
+        ),
+        encoding="utf-8",
+    )
+
+    def load_dataset(name, **_kwargs):
+        if name == "good":
+            return [{"prompt": "q", "response": "a"}]
+        return [{"prompt": "q", "response": "a"}, {"prompt": "later-row-is-invalid"}]
+
+    with pytest.raises(ValueError, match=r"source=bad.*row 1 missing required SFT field"):
+        train_cli._load_mixed_dataset_for_training(
+            str(manifest),
+            model_hub="hf",
+            dataset_loader_fn=None,
+            load_dataset=load_dataset,
+            load_from_disk=lambda _path: None,
+        )
+
+
+@pytest.mark.parametrize(
+    "row",
+    [
+        {"tokens": [1, 2], "prompt_mask": [True, False]},
+        {"image_base64": "fixture", "response": "answer"},
+        {"images_base64": ["fixture"], "response": "answer"},
+    ],
+)
+def test_mix_validation_accepts_current_sft_row_contracts(row):
+    train_cli._validate_sft_mix_source("source", [row])
+
+
+def test_mix_imports_shared_loader_only_once(tmp_path):
+    marker = tmp_path / "imports.txt"
+    loader = tmp_path / "loader.py"
+    loader.write_text(
+        f"from pathlib import Path\nwith Path({str(marker)!r}).open('a') as marker:\n"
+        "    marker.write('imported\\n')\n"
+        "def load_training_dataset(dataset_path, *, default_loader, **kwargs):\n"
+        "    return default_loader(dataset_path)\n",
+        encoding="utf-8",
+    )
+    manifest = tmp_path / "mix.json"
+    manifest.write_text(
+        json.dumps(
+            {
+                "version": 1,
+                "seed": 1,
+                "exhaustion": "renormalize",
+                "sources": [
+                    {"name": "first", "path": "first", "weight": 1},
+                    {"name": "second", "path": "second", "weight": 1},
+                ],
+            }
+        ),
+        encoding="utf-8",
+    )
+
+    train_cli._load_mixed_dataset_for_training(
+        str(manifest),
+        model_hub="hf",
+        dataset_loader_fn=str(loader),
+        load_dataset=lambda *_args, **_kwargs: [{"prompt": "q", "response": "a"}],
+        load_from_disk=lambda _path: None,
+    )
+
+    assert marker.read_text(encoding="utf-8").splitlines() == ["imported"]
+
+
+def test_mix_artifact_is_structured_and_sample_free(tmp_path):
+    dataset = WeightedMixedDataset(
+        [
+            DatasetMixSource("first", [{"prompt": "secret-one", "response": "answer-one"}], 1.0),
+            DatasetMixSource("second", [{"prompt": "secret-two", "response": "answer-two"}], 1.0),
+        ],
+        seed=1,
+        exhaustion="renormalize",
+    )
+
+    artifact_path = train_cli._write_dataset_mix_artifact(dataset, str(tmp_path))
+
+    assert artifact_path is not None
+    assert artifact_path.name.endswith(".epoch-0.json")
+    artifact_text = artifact_path.read_text(encoding="utf-8")
+    artifact = json.loads(artifact_text)
+    assert artifact["policy"] == "renormalize"
+    assert artifact["planned_rows"] == 2
+    assert "schedule_hash" in artifact
+    assert "secret-one" not in artifact_text
+    assert "answer-two" not in artifact_text
diff --git a/tests/test_mlx_training_cpu.py b/tests/test_mlx_training_cpu.py
index b9b81ee4..d9a3cccb 100644
--- a/tests/test_mlx_training_cpu.py
+++ b/tests/test_mlx_training_cpu.py
@@ -285,3 +285,63 @@ def test_adam8bit_matches_bias_corrected_adamw_for_uniform_moments():
         mx.eval(reference_model.parameters(), quantized_model.parameters())
 
     assert bool(mx.allclose(reference_model.weight, quantized_model.weight, atol=2e-5, rtol=2e-5).item())
+
+
+def test_adam8bit_dynamic_codebooks_match_cuda_reference():
+    mx = pytest.importorskip("mlx.core")
+
+    from areno.api.backend.mlx.optimizer import _mlx_dynamic_codebook
+    from areno.engine.optim.dynamic_quant import SIGNED_DYNAMIC_MAP, UNSIGNED_DYNAMIC_MAP
+
+    _require_mlx_device(mx)
+    signed = _mlx_dynamic_codebook(signed=True)
+    unsigned = _mlx_dynamic_codebook(signed=False)
+    mx.eval(signed, unsigned)
+
+    np.testing.assert_array_equal(np.array(signed), np.asarray(SIGNED_DYNAMIC_MAP, dtype=np.float32))
+    np.testing.assert_array_equal(np.array(unsigned), np.asarray(UNSIGNED_DYNAMIC_MAP, dtype=np.float32))
+
+
+def test_adam8bit_mlx_precision_callback_keeps_fp32_moments():
+    mx = pytest.importorskip("mlx.core")
+    nn = pytest.importorskip("mlx.nn")
+    from mlx.utils import tree_flatten, tree_unflatten
+
+    from areno.api.backend.mlx.optimizer import _quantized_adamw_class, apply_optimizer_update
+
+    _require_mlx_device(mx)
+    model = nn.Linear(8, 2, bias=False)
+    optimizer = _quantized_adamw_class()(
+        learning_rate=1e-3,
+        weight_decay=0.0,
+        state_precision_for_parameter=lambda _path, _parameter: "fp32",
+    )
+    path = tree_flatten(model.trainable_parameters())[0][0]
+    gradient = mx.ones_like(model.weight)
+    apply_optimizer_update(model, optimizer, tree_unflatten([(path, gradient)]))
+    state_names = {name for name, _ in tree_flatten(optimizer.state)}
+
+    assert any(name.endswith("m") for name in state_names)
+    assert any(name.endswith("v") for name in state_names)
+    assert not any(name.endswith(("m_q", "v_q", "m_scale", "v_scale")) for name in state_names)
+
+
+def test_mlx_provider_routes_embedding_by_identity_without_path_matching():
+    from areno.api.backend.mlx.provider import MlxModelProvider
+
+    embedding_weight = object()
+    per_layer_weight = object()
+    ordinary_weight = object()
+    embedding = SimpleNamespace(weight=embedding_weight)
+    per_layer_embedding = SimpleNamespace(weight=per_layer_weight)
+    model = SimpleNamespace(
+        model=SimpleNamespace(
+            embed_tokens=embedding,
+            embed_tokens_per_layer=[per_layer_embedding],
+        )
+    )
+    provider = MlxModelProvider(model, tokenizer=None, processor=None, config={})
+
+    assert provider.optimizer_state_precision("unexpected.path", embedding_weight) == "fp32"
+    assert provider.optimizer_state_precision("another.unexpected.path", per_layer_weight) == "fp32"
+    assert provider.optimizer_state_precision("embed_tokens.lookalike", ordinary_weight) == "8bit"
diff --git a/tests/test_multimodal_optimizer_cpu.py b/tests/test_multimodal_optimizer_cpu.py
index 6f61afda..6c7260c4 100644
--- a/tests/test_multimodal_optimizer_cpu.py
+++ b/tests/test_multimodal_optimizer_cpu.py
@@ -48,6 +48,11 @@ def test_cuda_adam8bit_matches_fp32_master_bias_corrected_updates():
         betas=kwargs["betas"],
         weight_decay=kwargs["weight_decay"],
     )
+    # Dynamic signed quantization deliberately uses the paper's asymmetric
+    # codebook, whose negative endpoint is -0.99296875 rather than -1.0.
+    # Accumulated updates therefore track FP32 within a small quantization
+    # budget instead of being bit-exact.
+    quantization_atol = 5e-6
 
     for gradient in (0.25, -0.5, 0.125, 1.0):
         reference_param.grad = torch.tensor([gradient])
@@ -56,5 +61,5 @@ def test_cuda_adam8bit_matches_fp32_master_bias_corrected_updates():
         reference.step()
         quantized.step()
         torch_reference.step()
-        torch.testing.assert_close(quantized_param, reference_param, atol=1e-6, rtol=1e-6)
-        torch.testing.assert_close(quantized_param, torch_param, atol=1e-6, rtol=1e-6)
+        torch.testing.assert_close(quantized_param, reference_param, atol=quantization_atol, rtol=1e-6)
+        torch.testing.assert_close(quantized_param, torch_param, atol=quantization_atol, rtol=1e-6)
diff --git a/tests/test_train_cli_config_cpu.py b/tests/test_train_cli_config_cpu.py
index ffd57030..7d28492d 100644
--- a/tests/test_train_cli_config_cpu.py
+++ b/tests/test_train_cli_config_cpu.py
@@ -26,10 +26,147 @@ def test_train_config_requires_ckpt():
 
 
 def test_train_config_requires_dataset_path():
-    with pytest.raises(UsageError, match="--dataset-path is required"):
+    with pytest.raises(UsageError, match="one of --dataset-path, --dataset-mix-config"):
         _trainer_config_from_options(**_options(dataset_path=None, algo="sft"))
 
 
+def test_train_config_accepts_dataset_mix_manifest_for_sft(tmp_path):
+    loader = tmp_path / "loader.py"
+    loader.write_text("def load_training_dataset(dataset_path, **kwargs):\n    return []\n", encoding="utf-8")
+    manifest = tmp_path / "mix.json"
+    manifest.write_text(
+        json.dumps(
+            {
+                "version": 1,
+                "seed": 42,
+                "exhaustion": "renormalize",
+                "sources": [
+                    {"name": "a", "path": "a.jsonl", "weight": 0.7},
+                    {"name": "b", "path": "b.jsonl", "weight": 0.3},
+                ],
+            }
+        ),
+        encoding="utf-8",
+    )
+
+    config = _trainer_config_from_options(
+        **_options(
+            algo="sft",
+            dataset_path=None,
+            dataset_mix_config=str(manifest),
+            dataset_loader_fn=str(loader),
+        )
+    )
+
+    assert config.dataset_path is None
+    assert config.dataset_mix_config == str(manifest)
+
+
+def test_train_config_accepts_repeated_dataset_sources_for_sft():
+    config = _trainer_config_from_options(
+        **_options(
+            algo="sft",
+            dataset_path=None,
+            dataset_sources=(
+                "alpaca_cleaned=yahma/alpaca-cleaned:0.7",
+                "stanford_alpaca=tatsu-lab/alpaca:0.3",
+            ),
+            dataset_mix_seed=7,
+            dataset_mix_exhaustion="cycle",
+            dataset_mix_samples_per_epoch=1000,
+        )
+    )
+
+    assert config.dataset_path is None
+    assert config.dataset_mix_config is None
+    assert config.dataset_sources == (
+        "alpaca_cleaned=yahma/alpaca-cleaned:0.7",
+        "stanford_alpaca=tatsu-lab/alpaca:0.3",
+    )
+    assert config.dataset_mix_seed == 7
+    assert config.dataset_mix_exhaustion == "cycle"
+    assert config.dataset_mix_samples_per_epoch == 1000
+
+
+def test_train_config_requires_at_least_two_dataset_sources():
+    with pytest.raises(UsageError, match="repeat the option at least twice"):
+        _trainer_config_from_options(
+            **_options(
+                algo="sft",
+                dataset_path=None,
+                dataset_sources=("only=repo/dataset:1",),
+            )
+        )
+
+
+def test_train_config_rejects_dataset_source_with_other_dataset_input(tmp_path):
+    manifest = tmp_path / "mix.json"
+    manifest.write_text("{}", encoding="utf-8")
+
+    with pytest.raises(UsageError, match="mutually exclusive"):
+        _trainer_config_from_options(
+            **_options(
+                algo="sft",
+                dataset_path=None,
+                dataset_mix_config=str(manifest),
+                dataset_sources=("first=one:1", "second=two:1"),
+            )
+        )
+
+
+def test_train_config_rejects_dataset_path_with_mix_manifest(tmp_path):
+    manifest = tmp_path / "mix.json"
+    manifest.write_text("{}", encoding="utf-8")
+
+    with pytest.raises(UsageError, match="mutually exclusive"):
+        _trainer_config_from_options(**_options(algo="sft", dataset_path="dataset", dataset_mix_config=str(manifest)))
+
+
+def test_train_config_rejects_mix_manifest_for_non_sft(tmp_path):
+    manifest = tmp_path / "mix.json"
+    manifest.write_text(
+        json.dumps(
+            {
+                "version": 1,
+                "seed": 1,
+                "exhaustion": "stop",
+                "sources": [
+                    {"name": "a", "path": "a", "weight": 1},
+                    {"name": "b", "path": "b", "weight": 1},
+                ],
+            }
+        ),
+        encoding="utf-8",
+    )
+
+    with pytest.raises(UsageError, match="currently supports --algo sft only"):
+        _trainer_config_from_options(**_options(algo="gspo", dataset_path=None, dataset_mix_config=str(manifest)))
+
+
+def test_sdk_config_enforces_dataset_mix_contract():
+    with pytest.raises(ValueError, match="mutually exclusive"):
+        TrainerConfig(
+            algo="sft",
+            ckpt="actor",
+            dataset_path="dataset",
+            dataset_mix_config="mix.json",
+        )
+    with pytest.raises(ValueError, match="currently supports algo='sft' only"):
+        PolicyTrainerConfig(
+            algo="gspo",
+            ckpt="actor",
+            dataset_path=None,
+            dataset_mix_config="mix.json",
+        )
+
+
+def test_dataset_mix_config_does_not_shift_existing_positional_config_arguments():
+    config = TrainerConfig("sft", "actor", "dataset", "cuda", "hf", None, "save")
+
+    assert config.save_path == "save"
+    assert config.dataset_mix_config is None
+
+
 def test_train_config_validates_model_hub():
     with pytest.raises(UsageError, match="--model-hub must be one of: hf, modelscope"):
         _trainer_config_from_options(**_options(algo="sft", reward_fn_path=None, reward_ckpt=None, model_hub="bogus"))
@@ -722,6 +859,47 @@ def fake_run(config):
     assert events == [("run", "sft")]
 
 
+def test_train_command_accepts_repeated_dataset_source_options(monkeypatch):
+    captured = []
+    monkeypatch.setattr(train_cli, "run", captured.append)
+
+    result = CliRunner().invoke(
+        train_cli.train_command,
+        [
+            "--algo",
+            "sft",
+            "--ckpt",
+            "actor",
+            "--model-hub",
+            "hf",
+            "--dataset-source",
+            "alpaca_cleaned=yahma/alpaca-cleaned:0.7",
+            "--dataset-source",
+            "stanford_alpaca=tatsu-lab/alpaca:0.3",
+            "--dataset-loader-fn",
+            "examples/sft/alpaca/dataset_loader.py",
+            "--dataset-mix-seed",
+            "7",
+            "--dataset-mix-samples-per-epoch",
+            "1000",
+            "--world-size",
+            "1",
+            "--tp-size",
+            "1",
+        ],
+    )
+
+    assert result.exit_code == 0, result.output
+    assert "dataset_mix" in unstyle(result.output)
+    assert "2 command-line sources" in unstyle(result.output)
+    assert captured[0].dataset_sources == (
+        "alpaca_cleaned=yahma/alpaca-cleaned:0.7",
+        "stanford_alpaca=tatsu-lab/alpaca:0.3",
+    )
+    assert captured[0].dataset_mix_seed == 7
+    assert captured[0].dataset_mix_samples_per_epoch == 1000
+
+
 def test_train_command_tunes_params_before_summary_and_run(monkeypatch):
     from areno.cli.auto_tune import AutoTuneCandidate, AutoTuneMeasurement, AutoTuneResult
 
diff --git a/tests/test_trainer_dataset_utils_cpu.py b/tests/test_trainer_dataset_utils_cpu.py
index 45368d1a..6431ccbf 100644
--- a/tests/test_trainer_dataset_utils_cpu.py
+++ b/tests/test_trainer_dataset_utils_cpu.py
@@ -1,9 +1,13 @@
 from __future__ import annotations
 
+import json
+import tempfile
 import unittest
+from pathlib import Path
 from types import SimpleNamespace
 
 from areno.api import data_utils
+from areno.api.data import DatasetMixSource, WeightedMixedDataset
 from areno.api.trainers import dpo as dpo_mod
 from areno.api.trainers import sft as sft_mod
 
@@ -32,6 +36,7 @@ class FakeSFTBackend:
     def __init__(self):
         self.closed = False
         self.train_calls = 0
+        self.train_rows = 0
 
     def init(self):
         return None
@@ -45,12 +50,32 @@ def get_tokenizer(self):
     def get_processor(self):
         return None
 
-    def train(self, _batch, _loss_fn, *, mini_bs, gradient_accumulation_steps):
+    def train(self, batch, _loss_fn, *, mini_bs, gradient_accumulation_steps):
         del mini_bs, gradient_accumulation_steps
         self.train_calls += 1
+        self.train_rows += len(batch)
         return {}
 
 
+class EpochAwareDataset:
+    def __init__(self):
+        self.epochs = []
+        self.summary_epochs = []
+
+    def set_epoch(self, epoch):
+        self.epochs.append(epoch)
+
+    def summary(self):
+        self.summary_epochs.append(self.epochs[-1])
+        return {"epoch": self.epochs[-1]}
+
+    def __len__(self):
+        return 1
+
+    def __getitem__(self, _index):
+        return {"prompt": "q", "response": "a"}
+
+
 def _sft_config(**overrides):
     """Return the minimal config shape SFTTrainer reads in CPU tests."""
 
@@ -58,11 +83,13 @@ def _sft_config(**overrides):
         "batch_size": 2,
         "epochs": 1,
         "gradient_accumulation_steps": 1,
+        "max_steps": None,
         "max_new_tokens": 2,
         "max_prompt_tokens": 2,
         "mini_bs": 1,
         "save_interval": 1,
         "save_path": None,
+        "metrics_log_dir": None,
     }
     defaults.update(overrides)
     return SimpleNamespace(**defaults)
@@ -180,6 +207,131 @@ def test_sft_fit_raises_when_all_rows_are_filtered(self):
         self.assertEqual(backend.train_calls, 0)
         self.assertTrue(backend.closed)
 
+    def test_sft_sets_and_reports_dataset_epoch_before_each_pass(self):
+        backend = FakeSFTBackend()
+        dataset = EpochAwareDataset()
+        trainer = sft_mod.SFTTrainer(
+            _sft_config(epochs=2),
+            instance=backend,
+            dataset=dataset,
+            reward_fn=None,
+            loss_fn=lambda _pack, _logprobs: None,
+        )
+
+        trainer.fit()
+
+        self.assertEqual(dataset.epochs, [0, 1])
+        self.assertEqual(dataset.summary_epochs, [0, 1])
+        self.assertEqual(backend.train_calls, 2)
+
+    def test_sft_writes_one_dataset_mix_plan_for_each_epoch(self):
+        backend = FakeSFTBackend()
+        dataset = WeightedMixedDataset(
+            [
+                DatasetMixSource("math", [{"prompt": "q", "response": "a"}], 0.7),
+                DatasetMixSource("code", [{"prompt": "r", "response": "b"}], 0.3),
+            ],
+            seed=42,
+            exhaustion="renormalize",
+        )
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            trainer = sft_mod.SFTTrainer(
+                _sft_config(epochs=2, metrics_log_dir=temp_dir),
+                instance=backend,
+                dataset=dataset,
+                reward_fn=None,
+                loss_fn=lambda _pack, _logprobs: None,
+            )
+
+            trainer.fit()
+
+            artifacts = sorted(Path(temp_dir).glob("dataset_mix_plan.*.epoch-*.json"))
+            self.assertEqual([path.name.rsplit(".epoch-", 1)[1] for path in artifacts], ["0.json", "1.json"])
+            summaries = [json.loads(path.read_text(encoding="utf-8")) for path in artifacts]
+
+        self.assertEqual([summary["epoch"] for summary in summaries], [0, 1])
+        for summary in summaries:
+            dataset.set_epoch(summary["epoch"])
+            self.assertEqual(summary["schedule_hash"], dataset.summary()["schedule_hash"])
+
+    def test_sft_trains_two_mixed_sources_end_to_end_with_cpu_backend(self):
+        backend = FakeSFTBackend()
+        dataset = WeightedMixedDataset(
+            [
+                DatasetMixSource("math", [{"prompt": "q", "response": "a"}], 0.7),
+                DatasetMixSource("code", [{"prompt": "r", "response": "b"}], 0.3),
+            ],
+            seed=42,
+            exhaustion="renormalize",
+        )
+        trainer = sft_mod.SFTTrainer(
+            _sft_config(),
+            instance=backend,
+            dataset=dataset,
+            reward_fn=None,
+            loss_fn=lambda _pack, _logprobs: None,
+        )
+
+        with self.assertLogs("areno.api.trainers.sft.SFTTrainer", level="INFO") as logs:
+            trainer.fit()
+
+        self.assertEqual(backend.train_calls, 1)
+        self.assertEqual(backend.train_rows, 2)
+        self.assertTrue(backend.closed)
+        progress_log = next(message for message in logs.output if "stage=dataset_mix_progress" in message)
+        self.assertIn("'rows_scheduled': 2", progress_log)
+        self.assertIn("'rows_filtered': 0", progress_log)
+        self.assertIn("'rows_trained': 2", progress_log)
+        self.assertIn("'target_tokens_trained': 2", progress_log)
+        self.assertIn("'name': 'math'", progress_log)
+        self.assertIn("'name': 'code'", progress_log)
+        self.assertIn("'observed_sample_proportion': 0.5", progress_log)
+        self.assertIn("'observed_token_proportion': 0.5", progress_log)
+
+    def test_sft_mix_progress_exposes_post_tokenization_filter_drift(self):
+        backend = FakeSFTBackend()
+        dataset = WeightedMixedDataset(
+            [
+                DatasetMixSource("valid", [{"prompt": "q", "response": "a"}], 0.5),
+                DatasetMixSource("empty", [{"prompt": "r", "response": ""}], 0.5),
+            ],
+            seed=42,
+            exhaustion="renormalize",
+        )
+        trainer = sft_mod.SFTTrainer(
+            _sft_config(),
+            instance=backend,
+            dataset=dataset,
+            reward_fn=None,
+            loss_fn=lambda _pack, _logprobs: None,
+        )
+
+        with self.assertLogs("areno.api.trainers.sft.SFTTrainer", level="INFO") as logs:
+            trainer.fit()
+
+        epoch_end = next(message for message in logs.output if "stage=dataset_mix_epoch_end" in message)
+        self.assertIn("'rows_scheduled': 2", epoch_end)
+        self.assertIn("'rows_filtered': 1", epoch_end)
+        self.assertIn("'rows_trained': 1", epoch_end)
+        self.assertIn("'name': 'empty'", epoch_end)
+        self.assertIn("'observed_sample_proportion': 0.0", epoch_end)
+
+    def test_sft_target_token_count_honors_encoded_loss_mask(self):
+        sequence = sft_mod._record_to_train_sequence(
+            {
+                "tokens": [1, 2, 3, 4],
+                "prompt_mask": [True, False, False, False],
+                "loss_mask": [False, False, True, False],
+            },
+            FakeTextTokenizer(),
+            max_prompt_tokens=4,
+            max_new_tokens=4,
+        )
+
+        self.assertIsNotNone(sequence)
+        self.assertEqual(sft_mod._sft_target_token_count(sequence), 1)
+
     def test_dpo_requires_explicit_prompt_chosen_rejected_schema(self):
         """DPO rows should not guess preference or prompt field aliases."""
         tokenizer = FakeTextTokenizer()