Skip to content

Add DASC recurrent state sparsity policy - #2409

Draft
kaix-nv wants to merge 3 commits into
mainfrom
kaix/dasc-state-sparsity
Draft

kaix-nv wants to merge 3 commits into
mainfrom
kaix/dasc-state-sparsity

Conversation

@kaix-nv

@kaix-nv kaix-nv commented Sep 11, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

Type of change: new feature

Adds an experimental modelopt.torch.sparsity.state_sparsity API for the first ModelOpt DASC milestone:

  • derives deterministic GDN whole-head decay horizons from A_log and dt_bias in CPU FP64;
  • selects the largest arbitrary positive Wmax that passes caller-supplied quality, lifecycle, and physical-storage evidence;
  • keeps DASC-NR zero recovery and DASC-WR suffix replay as explicit contracts;
  • records model, checkpoint, calibration provenance, masks, gates, measurements, and decay-storage dtype in versioned JSON-safe ModelOpt state;
  • validates model structure and storage-aware decay bounds during restore and deployment export;
  • fails closed for distributed Megatron GDN policies whose TP or PP geometry would be rank-local.

This change is policy-only and does not alter model execution. Runtime ragged checkpoint packing, suffix replay, KDA support, quantization, and GDN or KDA kernels are intentionally out of scope. A serving implementation must preserve convolution state and materialize ordinary dense recurrent state before continuation.

This draft supersedes #2375 and includes all review fixes accumulated there.

Usage

Import modelopt.torch.sparsity.state_sparsity, call calibrate with the model, configuration, and paired measurements, then call export_policy to obtain the deployment policy. The complete executable example and schema are in the sparsity guide.

Measurements contain paired dense-versus-DASC quality and physical checkpoint-storage results for every configured candidate.

Testing

  • Focused DASC suite: 60 passed, 1 optional Megatron skip
  • State, weight, and attention sparsity compatibility suites: 338 passed, 1 optional Megatron skip
  • Real Transformers Qwen3NextGatedDeltaNet smoke: passed on the prior reviewed tree
  • Pre-commit on all changed files: passed
  • Git diff check: clean

Before your PR is Ready for review

  • Is this change backward compatible?: yes
  • Copied code or new dependency?: no
  • New necessary tests?: yes
  • Changelog updated?: yes
  • Claude approval?: pending; this PR intentionally remains draft

Additional Information

Algorithm reference: https://arxiv.org/abs/2608.30386

The mode stores calibration evidence rather than treating weight-derived horizons as sufficient proof. A future runtime milestone should consume this policy in the serving prefix-cache serializer and loader and qualify matched-HBM capacity plus prefix-hit latency.

Signed-off-by: Kai Xu <kaix@nvidia.com>
## Summary

Consolidates the complete reviewed fix stack for #2375 into one DCO-safe
commit:

- harden package exports, measurement semantics, wrapper handling, and
actionable calibration errors
- validate exact installed GDN identities plus ModelOpt dynamic
subclasses; reject lookalikes, ordinary subclasses, incomplete layers,
and partial layer sets
- make stale checkpoints saveable and restorable while keeping
deployment export strict
- make DASC recalibration replace and deduplicate existing mode state
without stale-metadata refresh
- record the declared decay-parameter checkpoint storage dtype and use
derived FP16/BF16/FP32 rounding bounds
- preserve BF16/FP16 storage and wider/cross-dtype reload compatibility
without globally widening FP32 tolerance
- add installed Transformers path coverage, optional Megatron gating,
lifecycle, tamper, lossy-cast, and mixed-layer regressions
- document the explicit storage-dtype contract

This consolidated PR supersedes the mechanically stacked review-fix PRs
#2377, #2378, #2379, #2380, #2382, #2383, #2384, and #2385. Its tree is
byte-identical to the independently reviewed leaf commit from #2386.

## Validation

- focused DASC suite: 23 passed, 1 absent optional Megatron skip
- DASC plus weight sparsity plus attention sparsity compatibility suite:
134 passed, 1 optional skip
- DASC package coverage: 408/408 statements, 100%
- full pre-commit on all touched files: passed
- real Transformers Qwen3NextGatedDeltaNet BF16 storage to FP32 reload
smoke: passed
- commit author and Signed-off-by identity both use kaix-nv
<kaix@nvidia.com>

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **New Features**
* Added support for configuring decay-parameter storage precision with
FP16, BF16, or FP32.
  * Added safer recalibration that replaces existing DASC state.
* Expanded compatibility with supported GDN adapter classes and model
wrappers.
* Added improved validation for sparsity policies, measurements, model
structure, and decay parameters.
  * Added support for perplexity-retention values above 1.

* **Documentation**
* Clarified evaluation responsibilities, recalibration behavior,
stale-policy handling, supported adapters, and dtype requirements.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: kaix-nv <kaix@nvidia.com>
Signed-off-by: Kai Xu <kaix@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Sep 11, 2026

Copy link
Copy Markdown

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@coderabbitai

coderabbitai Bot commented Sep 11, 2026

Copy link
Copy Markdown
Contributor

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

  • Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

Comment @coderabbitai help to get the list of available commands.

@github-actions

Copy link
Copy Markdown
Contributor
PR Preview Action v1.8.1

QR code for preview link

🚀 View preview at
https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-2409/

Built to branch gh-pages at 2026-09-11 18:41 UTC.
Preview will be ready when the GitHub Pages deployment is complete.

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.

1 participant