PyTorch ao
To use experimental optimizers (AdamWFp8, AdamW4bit, AdamW8bit) from Pytorch Ao, please install the package as shown below.
Some experimental optimizers are already present in regular Pytorch, so please re-check if you actually need this package!
Installation
Stable Release from the PyTorch index
pip install torchao --extra-index-url https://download.pytorch.org/whl/cu121 # full options are cpu/cu118/cu121/cu124Nightly release
pip install --pre torchao-nightly --index-url https://download.pytorch.org/whl/nightly/cu121 # full options are cpu/cu118/cu121/cu124Native NVFP4 LoRA with DDP
For a serialized, frozen native TorchAO weight-only NVFP4 base with ordinary PEFT LoRA, DDP automatically verifies component layouts and synchronizes the raw NVFP4 components from rank zero before wrapping the model. LoRA adapters continue through ordinary DDP reduction; no configuration flag is required. This path is qualified for two-rank regular weight-only LoRA training, checkpoint resume, and adapter save. Native dynamic activation NVFP4 is separately qualified for dense merge-aware LoRA on SM100+ hardware with unsharded runs, FSDP2, and DeepSpeed ZeRO-1/2/3, including adapter metadata and format-preserving export; see NVFP4 LoRA for constraints. These results do not establish dynamic DDP support.
Native NVFP4 LoRA with Tensor Parallelism
For static, serialized TorchAO NVFP4WeightOnlyConfig checkpoints, ordinary PEFT LoRA is prepared after adapter loading when tensor_parallel_size is greater than one. The qualified path is the Llama-style q_proj/o_proj colwise/rowwise layout on two TP ranks, validated with FP32 and BF16 compute. It preserves the frozen NVFP4 base forward, sums rowwise adapter deltas across the TP mesh, and reduces gradients only for replicated LoRA factors.
Use biasless LoRA factors with lora_dropout: 0. DoRA, adapter variants, mixed-batch adapter selection, gathered colwise outputs, dynamic NVFP4 activation quantization, and TP plans outside the supported colwise/rowwise layouts are rejected.