Skip to content

feat(runtime): compress checkpoint activations to FP8 - #538

Closed
xsuler wants to merge 4 commits into
mainfrom
feat/fp8-checkpoint-activations
Closed

feat(runtime): compress checkpoint activations to FP8#538
xsuler wants to merge 4 commits into
mainfrom
feat/fp8-checkpoint-activations

Conversation

@xsuler

@xsuler xsuler commented Sep 1, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • store BF16 tensors retained at activation-checkpoint boundaries as group-wise FP8 E4M3 by default on CUDA
  • restore BF16 tensors before backward recomputation while leaving the original forward unchanged
  • allow disabling compression with --no-fp8-ckpt-activations; activation checkpointing itself still controls whether the path is used
  • drop completed rollout state by default to reduce memory, with --keep-rollout-state as the explicit opt-out
  • align SDK/runtime defaults, CLI, Dashboard controls and generated commands, bundled WebUI assets, tests, and user documentation
  • support group-128 by default plus per-token/group-256 scaling, stochastic rounding, warm-up range collection, and per-layer BF16 fallback through the Python config
  • report boundary original/stored bytes and storage reduction in training metrics
  • enable Ling/Bailing V3 activation checkpointing for every attention block and dense MLP block while keeping sparse routing outside recomputation

Implementation

The compression hook wraps only the non-reentrant checkpoint call. PyTorch's checkpoint-owned saved-tensor hook continues to handle normal layer intermediates, while the outer hook packs the BF16 checkpoint input into FP8 plus FP32 scales. The packed payload does not retain the original BF16 tensor. Normalized MoE inputs and routing metadata stay outside the compressed boundary path.

On CUDA, FP8 boundary storage follows activation checkpointing and is enabled by default. --no-fp8-ckpt-activations retains BF16 boundary tensors. MLX leaves this CUDA-only feature disabled. Completed rollout KV/cache state is released by default on both backends; --keep-rollout-state restores the previous retention behavior.

The default CLI path uses deterministic E4M3 group-128 scaling. Advanced SDK fields allow per-token/group-256 scaling, stochastic rounding, warm-up steps, and explicit fallback layer indices.

Validation

  • pytest -q tests/test_recompute_cpu.py — 10 passed (including compiled-counter reset regression)
  • focused config/default propagation tests — 6 passed
  • focused CLI/default/help tests — 2 passed
  • Dashboard command-generation test — 1 passed
  • Dashboard production bundle rebuilt successfully
  • pre-commit run --all-files --hook-stage manual — passed

CPU tests cover storage reduction, BF16 restoration, exact original-forward output, gradient cosine similarity >= 0.99, RNG isolation, warm-up/fallback behavior, default/opt-out semantics, WebUI command generation, and CLI/runtime propagation.

Full CUDA long-sequence peak-memory, throughput, and three-seed training/evaluation acceptance benchmarks require GPU infrastructure and are not reported by this local change.

Closes #534

@xsuler
xsuler deployed to push_check September 1, 2026 06:59 — with GitHub Actions Active
@xsuler
xsuler deployed to push_check September 1, 2026 07:13 — with GitHub Actions Active
@xsuler
xsuler deployed to push_check September 1, 2026 07:29 — with GitHub Actions Active
@xsuler xsuler closed this Sep 1, 2026
@xsuler
xsuler deleted the feat/fp8-checkpoint-activations branch September 3, 2026 03:33
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Compress activation-checkpoint boundary tensors to FP8

1 participant