Add an NVFP4 quantization converter - #3914
Conversation
|
|
fc8b67b to
2a36c34
Compare
| ngpu: int = 4 | ||
| disabled: bool = False | ||
| skip_rocm_test: bool = False | ||
| skip_if_no_blackwell: bool = False |
There was a problem hiding this comment.
torchtitan CI doesn't have blackwell
@syed-ahmed any plan to sponsor?
There was a problem hiding this comment.
Blackwell CI support is in the works. Once it's allocated in the pytorch org, we can start using it.
There was a problem hiding this comment.
Maybe remove this test for now? CI has no signal so can't guard it anyway.
There was a problem hiding this comment.
So remove all gpu tests?
There was a problem hiding this comment.
We run the torchtitan tests in our internal blackwell CI and will maintain the tests.
There was a problem hiding this comment.
Your CI has blackwell so you don't need this flag?
| @override( | ||
| "nvfp4_feed_forward", | ||
| target=FeedForward.Config, | ||
| fqns=["layers.*.feed_forward", "*.layers.*.feed_forward"], | ||
| description="NVFP4 sequence-parallel FFN block (fp4 all-gather, no bf16 gather).", | ||
| ) |
There was a problem hiding this comment.
override should be mostly used for unblocking work outside the repo. For in-repo quantization support, we should put in https://github.com/pytorch/torchtitan/tree/main/torchtitan/components/quantization
You could make a converter (as a convenient util) following other quantization applications https://github.com/pytorch/torchtitan/blob/main/torchtitan/components/quantization/float8.py#L53
There was a problem hiding this comment.
@tianyu-l I thought we were moving away from converters and writing quantized modules explicitly? Is there a change in guidance there?
There was a problem hiding this comment.
@syed-ahmed
It might be slightly subtle:
- We are moving away from "first create an
nn.Linearand then callquantize_()to convert the module". See. enable graph trainer + mxfp8 composability #3558 - We are moving towards "creating a quantized module directly", in torchtitan out of a quantized module config.
- How do you obtain this config is flexible, via converter / override / creating one directly.
- Here I'm saying we shouldn't use override if this is some code we seriously support in the repo. Converters are just config utils, which I don't see problem working with. E.g. for LoRA it's pretty convenient https://github.com/pytorch/torchtitan/blob/main/torchtitan/components/lora.py#L141
- I don't have problem if someone insists on not using converter, but then we need to introduce more entries in
model_registry()which themselves are model config utils. https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/qwen3/__init__.py#L606
| self.tp_group = tp_mesh.get_group("tp") | ||
| self.world_size = tp_mesh.size() | ||
| if self.tensor_parallel_style is None: | ||
| self.tensor_parallel_style = _infer_tp_style(self._sharding_config) |
There was a problem hiding this comment.
why is inferring needed here? is there precendent for doing it this way vs making the user do it?
There was a problem hiding this comment.
Why is inferring needed here?
Not for any normally-specified TorchTitan model. Every model sets weight placement via colwise_config()/rowwise_config(), and the two override factories then set tensor_parallel_style explicitly on top. Inference at 345-346 fires only when a config is under-specified (style left None with a populated sharding_config) — i.e. hand-construction, which no in-repo caller does.
Precedent for doing it this way vs making the user do it?
The implicit precedent is the weight placement — the repo's sole canonical colwise/rowwise encoding across every model. Instead of inferring, the override can trigger an assertion when the sharding configuration is incomplete.
|
|
||
|
|
||
| def _nvfp4_rowwise_sp(x_BLD, w_local, bias, sr_seed, sign_vector, tp_group, world_size): | ||
| """Rowwise NVFP4 over a full-sequence feature shard, returning a seq shard. |
There was a problem hiding this comment.
is "rowwise" here talking about the nvfp4 recipe (outer scale rowwise, inner scale 1x16), or something else?
There was a problem hiding this comment.
"rowwise" refers to rowwise sequence parallel layer. It isn't related to nvfp4 recipe.
can you provide some more context on all of these? |
2a36c34 to
b97cecb
Compare
Only the linears are replaced with NVFP4. The remaining gemms in MOE are unchanged.
The idea is do all-gather on the quantized NVFP4 tensors to save comm bandwidth. Stock TorchTitan does bf16 all-gather separately from The TorchAO NVFP4ColwiseParallel and NVFP4RowwiseParallel bake in comms to do NVFP4 All-Gather, the override needs to ensure TorchAO agrees with TorchTitan sharding scheme.
NVFP4 tracks stochastic rounding seed and sign vector for random hadamard transform in its checkpoint state. When you export to HuggingFace you lose them. The PR's behavior is to redraw a new SR seed and sign vector. cc: @vkuzo |
| def parallelize(self, parallel_dims: ParallelDims) -> None: | ||
| self._tp_active = parallel_dims.tp_enabled | ||
| if self._sharding_config is not None: | ||
| # Declare the runtime buffers (replicated) so _distribute_states | ||
| # and DCP handle them alongside weight/bias. | ||
| sc = self._sharding_config | ||
| self._sharding_config = replace( | ||
| sc, | ||
| state_shardings={ | ||
| **sc.state_shardings, | ||
| "_sr_seed": _replicated_layout(), | ||
| "_rht_sign_vector": _replicated_layout(), | ||
| }, | ||
| ) | ||
| self._tp_style = _infer_tp_style(self._sharding_config) | ||
| self._validate(parallel_dims) | ||
| self._cache_buffer_spec(parallel_dims) | ||
| super().parallelize(parallel_dims) |
There was a problem hiding this comment.
What are we trying to achieve by modifying this code? Some more context would be helpful.
There was a problem hiding this comment.
TorchAO's NVFP4Linear isn't a plain Linear. It carries two extra runtime buffers (_sr_seed, _rht_sign_vector) and a hard kernel constraint that each local GEMM dim must be a multiple of 128. TorchTitan's Module protocol requires every piece of state to be declared before it distributes params, and the stock colwise/rowwise sharding_config only knows about weight and bias. So we override parallelize() to do the NVFP4-specific setup, then delegate to the base via super().parallelize().
| ) | ||
| out_tp = Shard(-1) if self._tp_style == "colwise" else Partial() | ||
| mesh, placements = _swap_tp_placement(x, out_tp) | ||
| return DTensor.from_local(y, mesh, placements, run_check=False) |
There was a problem hiding this comment.
We are migrating from DTensor to spmd_types, could you make sure it (only) works with spmd_types backend and debug.spmd_typechecking?
There was a problem hiding this comment.
The spmd local_map-region pattern was modeled on RoutedExperts in torchtitan/models/common/moe/. It wraps an opaque local-compute op in a module-level spmd.local_map region via sharding_config.local_map.
| ngpu: int = 4 | ||
| disabled: bool = False | ||
| skip_rocm_test: bool = False | ||
| skip_if_no_blackwell: bool = False |
There was a problem hiding this comment.
Maybe remove this test for now? CI has no signal so can't guard it anyway.
| self._sr_seed = self._materialize_buffer( | ||
| torch.randint( | ||
| -(2**63), 2**63 - 1, (1,), dtype=torch.int64, device=dev | ||
| ) | ||
| ) | ||
| self._rht_sign_vector = self._materialize_buffer( | ||
| _make_rht_sign_vector(None, device=dev) | ||
| ) |
There was a problem hiding this comment.
I don't think you'd need to introduce _materialize_buffer -- I feel it can be configured right e.g. https://github.com/pytorch/torchtitan/blob/main/torchtitan/overrides/fused_swiglu.py#L500-L504
There was a problem hiding this comment.
RHT must be replicated across TP ranks because RHT transformation must be consistent along quantized dimension and gemm contraction dim. The broadcast is load-bearing and the spmd_types R path won't do it. I added comment to code.
| ngpu: int = 4 | ||
| disabled: bool = False | ||
| skip_rocm_test: bool = False | ||
| skip_if_no_blackwell: bool = False |
|
The
|
| ngpu: int = 4 | ||
| disabled: bool = False | ||
| skip_rocm_test: bool = False | ||
| skip_if_no_blackwell: bool = False |
There was a problem hiding this comment.
Your CI has blackwell so you don't need this flag?
d0ce31f to
db06503
Compare
Fold the native-checkpoint exclusion assertion into the exposes-buffers test and drop test_nvfp4_native_checkpoint_excludes_runtime_buffers. The dropped test's load_state_dict + torch.equal round-trip added no distinct failure mode: loading a weight-only state dict cannot touch the non-persistent runtime buffers by construction. The merged test keeps the real invariant -- a native checkpoint carries only the stock weight. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
bbb4492 to
5d2e9ff
Compare
4ab90ed to
3c1999e
Compare
3713874 to
3a299f2
Compare
|
@tianyu-l @vkuzo I fixed the lint errors and added
4xFSDP GB200 |
|
|
||
| 1. Train with NVFP4 for most of the run. | ||
| 2. When exact loss recovery matters, switch the linear-layer GEMM inputs to a higher precision shortly before learning-rate decay begins. | ||
| 3. Prefer switching only the forward-pass GEMMs to bf16 or, potentially, MXFP8. |
There was a problem hiding this comment.
can you explain this recommendation in more detail? Naively I would have expected the opposite, but tbh it's best to either not say this, or if we do say it it should be backed with data
There was a problem hiding this comment.
This also comes from Appendix D. Switching to Higher Precision from the same paper. I can drop the part about MXFP8, which wasn't tested in the paper.
There was a problem hiding this comment.
makes sense! how about the following to frame it:
- we quote https://arxiv.org/pdf/2509.25149 here
- we extract some of the things in Appendix D here in plain english and explain how do do those in titan (to save the reader the click to the paper + agent search titan code)
I think that would be a good combination of sharing and attributing the tips+tricks found by the NVIDIA researchers so far and also making it clear that these are not universal rules that are guaranteed to work. It's also good to link back to the paper for those who want to understand where these guidelines are coming from
There was a problem hiding this comment.
Prefer switching only the forward-pass GEMMs to bf16 or, potentially, MXFP8.
btw, both of these are not yet supported in titan + torchao, right? It should be doable and useful to implement, but if not supported yet would be good to clarify it in this doc!
There was a problem hiding this comment.
Yes, you can't switch only fwd to bf16 and keep nvfp4 bwd today. The simplest option is to switch all layers to bf16 at end of training. I'll make the wording precise about what is available today.
| ### NVFP4 Training Recommendations | ||
|
|
||
| 1. Train with NVFP4 for most of the run. | ||
| 2. When exact loss recovery matters, switch the linear-layer GEMM inputs to a higher precision shortly before learning-rate decay begins. |
There was a problem hiding this comment.
do we have data backing this?
There was a problem hiding this comment.
This comes from the Pretraining Large Language Models with NVFP4. See Appendix D. Switching to Higher Precision
# Conflicts: # torchtitan/models/llama3/config_registry.py
|
Merging since this PR doesn't touch the failing tests in CI. |
Merged upstream/main (merge ddb4173). No conflicts, no replay required. - bcc0929 NVFP4 quantization converter (pytorch#3914) - d905f73 dependabot: pypa/gh-action-pypi-publish 1.14.1 -> 1.14.2 (CI only) The protocol replays llama3/ changes onto ezpz/agpt/, so bcc0929 was checked line by line: it is purely additive (zero removed/modified lines in llama3/config_registry.py; +3 each in the shared quantization/__init__.py and utils.py with nothing existing changed). NVFP4 is additionally NVIDIA-only (Blackwell FP4, every GEMM dim divisible by 128), so it is inapplicable on XPU regardless. Nothing to port.



Summary
Add NVFP4 (NVIDIA Blackwell, sm_100+) as a first-class quantization recipe using the converter mechanism, alongside
Float8LinearConverterandMXFP8LinearConverter(#3558): the quantized module is built from a config node viamodel_registry(converters=...), so NVFP4 is selectable through the normal config path and tested in the standard flow.What changed
torchtitan/components/quantization/nvfp4.py:NVFP4LinearConverter(aLinear.Config->NVFP4Linear.Configleaf swap, mirroringMXFP8LinearConverter) andNVFP4Linear.NVFP4Linearreuses torchao's statefulNVFP4Linear(weight/bias, the_sr_seed/_rht_sign_vectorruntime buffers, RHT logic, state-dict handling) and adds:_sr_seedis a per-rank, non-persistent stochastic-rounding key (SR is unbiased and no quantized values cross the wire, so distinct seeds per rank are correct; a Philox key needs no checkpointing)._rht_sign_vectoris drawn per-rank at random but broadcast to a replicated, persistent buffer, because the Hadamard basis must match across TP ranks (rowwise TP shards the GEMM contraction dim, and the transform only cancels when both operands share the sign vector).spmd_typesbackend. torchao's functional NVFP4 op (nvfp4_mm_triton) is an opaque Triton autograd Function that DTensor cannot dispatch on, so TP runs it inside anspmd.local_mapregion (modeled onRoutedExperts): the framework converts activations to local shards on entry and re-types the output on exit, andnvfp4_mm_tritonis registered viaspmd.register_local_autograd_function.qwen3_8b_nvfp4_mixedandlama3_8b_nvfp4_mixedconfigs convert only the leading decoder layers to NVFP4 and keep thelast ceil(n_layers * 0.15)layers and lm_head in bf16.Design notes
spmd_types-only TP. The opaque NVFP4 op needs an explicit local-compute region to type-check under SPMD; the DTensor backend has no way to dispatch it. So TP is supported only under
spmd_backend='spmd_types', and is validated withdebug.spmd_typechecking(an eager type checker that confirms the local_map colwise/rowwise output and input-gradient types). FSDP-only / single-GPU runs are unaffected and take torchao's local functional path directly.bf16 collectives at the block boundary. The converter keeps the model's stock bf16 TP collectives and quantizes only the GEMM -- it does not move fp4 codes over the wire. Moving fp4 would only help the column-parallel all-gather (the rowwise reduction needs a real sum and stays bf16), scales only with TP degree x sequence length, and is zero at TP=1. For NVFP4's targets -- large-K MoE models trained low-TP / high-EP, where bf16 all-to-all dominates and MoE experts are out of scope here -- the fp4 compute win is the prize and is fully retained. Keeping bf16 collectives also lets
NVFP4Linearstay a leaf swap that composes with FSDP/TP through the standardsharding_config.Numerical and performance validation
Llama 3 8B was trained on C4 dataset for 763 steps (200,015,872 tokens), sequence length 2048, global batch size 128, 4-way FSDP, TP 1. NVFP4, MXFP8, BF16 used the same local batch size of 32 with global accumulation of 1.
Experiment Setup
Table 1 — Eager Trainer 200M-token run
Takeaway: NVFP4 has higher tokens per second (TPS) and lower peak memory usage than MXFP8 and BF16. It delivers ~1.45× throughput at ~0.59× memory with matched loss (1.274 vs 1.268) against BF16. MXFP8 is ~1.28 faster than BF16 but gives no memory relief because it saves BF16 activations for backwards pass. NVFP4 beats MXFP8 by ~1.13× on throughput while using only ~0.57× its memory (103 vs 180 GiB, ~77 GiB less). NVFP4 saves quantized activations and scale factors for backwards pass, reducing memory pressure.
Training Loss Convergence
Eager Trainer 200M-token run for NVFP4, MXFP8, BF16

AI tools used