Skip to content

Support FSDP inference with precomputed quantized weights #1822

Description

@Davids048

Summary

Fully Sharded Data Parallel (FSDP) inference represents model parameters as distributed tensors whose local values are shards of the logical weight. Precomputed MXFP8 and NVFP4 conversion must produce quantized weights and scales that remain correct across FSDP all-gather, resharding, and offload operations.

PR #1796 rejects precomputed transformer quantization with FSDP inference until that lifecycle is implemented. The allowlist is enforced in fastvideo/models/loader/fsdp_load.py.

Review context: #1796 (comment)

Behavior implemented by PR #1796

FSDP inference configuration Behavior
No transformer quantization Allowed
NVFP4QATTrainConfig Allowed; weights remain FSDP-managed and quantization occurs during execution
MXFP8Config Explicitly rejected
NVFP4Config Explicitly rejected
Other transformer quantization configurations Explicitly rejected
Quantized inference without FSDP Unchanged

Support it if we need to do training with FSDP, with quantized layers.

Required work

  • Quantize at a point where the complete logical weight is available, or define and validate a shard-local algorithm whose block boundaries and scales match whole-weight quantization.
  • Register quantized weights and scales with an FSDP-managed distributed representation.
  • Preserve quantized state across FSDP all-gather, resharding, CPU offload, and module materialization.
  • Avoid quantizing only weight.to_local() and then treating that local shard as a complete persistent inference weight.
  • Define how post-load transformations, including LoRA merging, trigger distributed requantization.

Tests

  • Add a multi-process parity test comparing FSDP quantized inference with single-device quantized inference.
  • Verify quantized weight and scale shapes on every rank.
  • Exercise all-gather and reshard cycles across repeated forward passes.
  • Cover CPU offload when FSDP inference supports it.
  • Verify that LoRA-merged weights are quantized consistently on every rank.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

scope: attentionAttention backends (VSA, STA, Flash, etc.)scope: distributedSP, FSDP, USP, multi-nodescope: inferenceInference pipeline, serving, CLIscope: trainingTraining pipeline, methods, configs

Type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions