[feat] 1/n Unified Quantization Framework

[kailash109]

Overview

This PR extends from #1940 and #1898 to propose a general unified framework for implementing quantization schemes in SkyRL -- the blockwise fp8 and mxfp8 implementations are refactored to fit in this implementation

follow-up PR implementing nvfp4 recipe on top of this: #1967

Additions

Added to skyrl/backends/skyrl_train/quantization:

• base.py: Introduce QuantizedModelLayout and QuantizationStrategy classes: the QuantizedModelLayout class defines a model-specific weight structure, providing model-type lookup/validation, a method for splitting packed Bridge tensors (ie. qwen gate/up expert weights), and expert export layout required from Mbridge. QuantizationStrategy defines a complete quantization spec, providing a should_quantize() method to select weight categories for quantization, methods to build the TE config, megatron provider, and runtime env for the particular quantization format, as well as a serialize_weight() functionality which emits the packed weights during weight syncs (this replaces the original logic of iter_serialized_fp8_tensors() during weight syncs).
• mxfp8.py/blockwise_fp8.py implement the persistent fp8 + mxfp8 training pipelines in this new unified scheme (the code is mostly the same, just passing the data to the respective *Strategy classes.
• model_layouts.py: Creates the same qwen3.5 + moe layouts from the previous PRs in the new scheme using QuantizedModelLayout
• megatron.py: builds Megatron quantization config given QuantizationStrategy
• vllm.py: Single method to rename checkpoint tensor to fused VLLM parameter/shard

These abstractions make it very easy to implement new quantization strategies with arbitrary packing/scale logic + targetting arbitary layers/modules of models. Rather than have to plumb through new configs for TE/Megatron and add specific per-layer quantization logic + additional weight serialization in the WeightExtractor, everything is automatically handled through the QuantizationStrategy class (and/or QuantizedModelLayout if adding quantization support for new model). This also enables significant reuse of existing quantization frameworks if supported models/vLLM weight recasting is already implemented.

Current WIP is implementing the nvfp4 recipe from humans& on top of these classes.

Validation

All CPU tests pass including new tests added to test emitted weight/scales from existing FP8/MXFP8 strategy + qwen3 layout.

e2e training run results:

[gemini-code-assist]

Code Review

This pull request introduces support for expert-only MXFP8 quantization on Blackwell GPUs (SM100/SM103) using Transformer Engine and vLLM, including new quantization strategies, model layouts, and benchmark scripts. Feedback on the changes highlights a potential RuntimeError in blockwise_cast_to_fp8 due to non-contiguous tensors, a missing copy import in the inference server utilities, and a misleading error message in the MoE weight splitting logic.