Skip to content

[feature] online quantize weights when loading weights #1500

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
haoyangli0109 wants to merge 1 commit into from
Closed

[feature] online quantize weights when loading weights #1500

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
334 changes: 333 additions & 1 deletion atom/model_loader/loader.py
View file
Open in desktop

Large diffs are not rendered by default.

55 changes: 48 additions & 7 deletions atom/model_ops/linear.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,11 @@
from aiter.tuned_gemm import tgemm
from aiter.utility import fp4_utils
from atom.config import QuantizationConfig, get_current_atom_config
from atom.quant_spec import LayerQuantConfig, should_skip_online_quant
from atom.quant_spec import (
LayerQuantConfig,
should_skip_online_quant,
should_stream_online_quant,
)
from atom.model_ops.utils import (
atom_parameter,
normalize_e4m3fn_to_e4m3fnuz,
Expand Down Expand Up @@ -414,21 +418,46 @@ def __init__(
divide(s, self.tp_size) for s in self.output_partition_sizes
]

# Decide whether this layer streams its online quantization: quantize it
# right after its weights finish loading (freeing the source BF16),
# instead of loading the whole model then quantizing. Same decision as
# FusedMoE (see quant_spec.should_stream_online_quant). When enabled the
# weight is allocated on the meta device and the loader materializes it
# on first touch (see model_loader/loader.py).
self._stream_online = self.source_quant_dtype is None and (
should_stream_online_quant(quant_config, prefix, quant_type, params_dtype)
)
# Capture the intended device now: the model is constructed under
# `set_default_device(self.device)`, but by load time the default is
# reset, so we must remember where to materialize.
self._load_device = torch.empty(0).device if self._stream_online else None
# Single switch for every parameter allocated below. A streaming layer
# starts entirely on meta -- weight, bias and scales alike -- so it costs
# no real memory until the loader materializes each param on first touch;
# `None` means "default device", i.e. the classic behaviour.
param_device = "meta" if self._stream_online else None

if self.source_quant_dtype is not None:
weight_size = (self.output_size, self.input_size)
self.weight = atom_parameter(
torch.empty(weight_size, dtype=self.source_quant_dtype)
torch.empty(
weight_size, dtype=self.source_quant_dtype, device=param_device
)
)
else:
weight_size = (
(self.output_size, self.input_size)
if params_dtype not in [dtypes.fp4x2, dtypes.i4x2]
else (self.output_size, self.input_size // 2)
)
self.weight = atom_parameter(torch.empty(weight_size, dtype=params_dtype))
self.weight = atom_parameter(
torch.empty(weight_size, dtype=params_dtype, device=param_device)
)
if bias:
output_type = get_current_atom_config().torch_dtype
self.bias = atom_parameter(torch.empty(self.output_size, dtype=output_type))
self.bias = atom_parameter(
torch.empty(self.output_size, dtype=output_type, device=param_device)
)
self.bias.weight_loader_process = self.weight_loader_process
else:
self.register_parameter("bias", None)
Expand All @@ -438,19 +467,29 @@ def __init__(
if quant_type != QuantType.No and self.source_quant_dtype is None:
if quant_type == QuantType.per_Tensor:
self.weight_scale = atom_parameter(
torch.empty(len(self.output_partition_sizes), 1, dtype=dtypes.fp32)
torch.empty(
len(self.output_partition_sizes),
1,
dtype=dtypes.fp32,
device=param_device,
)
)
if not layer_quant_config.is_dynamic:
self.input_scale = atom_parameter(
torch.empty(
len(self.output_partition_sizes), 1, dtype=dtypes.fp32
len(self.output_partition_sizes),
1,
dtype=dtypes.fp32,
device=param_device,
)
)
self.input_scale.weight_loader_process = self.weight_loader_process
self.input_scale.weight_loader = self.weight_loader
elif quant_type == QuantType.per_Token:
self.weight_scale = atom_parameter(
torch.empty(self.output_size, 1, dtype=dtypes.fp32)
torch.empty(
self.output_size, 1, dtype=dtypes.fp32, device=param_device
)
)
elif quant_type == QuantType.per_1x128:
scale_dtype = (
Expand All @@ -463,6 +502,7 @@ def __init__(
(self.output_size + 127) // 128,
(self.input_size + 127) // 128,
dtype=scale_dtype,
device=param_device,
)
)
elif quant_type == QuantType.per_1x32:
Expand All @@ -471,6 +511,7 @@ def __init__(
self.output_size,
(self.input_size + 31) // 32,
dtype=dtypes.fp8_e8m0,
device=param_device,
)
)
self.weight.weight_loader_process = self.weight_loader_process
Expand Down
119 changes: 101 additions & 18 deletions atom/model_ops/moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,11 @@
shuffle_weights,
)
from atom.plugin.vllm.moe import FusedMoEDecoratorForPluginMode
from atom.quant_spec import LayerQuantConfig, should_skip_online_quant
from atom.quant_spec import (
LayerQuantConfig,
should_skip_online_quant,
should_stream_online_quant,
)
from atom.quantization.quark.utils import (
dequant_weight_online,
quant_weight_online,
Expand Down Expand Up @@ -628,13 +632,22 @@ def create_weights(
params_dtype: torch.dtype,
**extra_weight_attrs,
):
# When this expert module streams its online quantization, allocate the
# source w13/w2 buffers on the meta device (0 real memory at build time);
# the loader materializes them on first touch and _online_quant frees
# them right after quantizing. Only the initial (unquantized-source)
# allocation is meta -- the target buffers built later in _online_quant
# go through the Fp8/Mxfp4 method, which allocates real tensors.
weight_device = "meta" if getattr(layer, "_stream_online", False) else None

# Fused gate_up_proj (column parallel)
w13_weight = atom_parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size,
dtype=params_dtype,
device=weight_device,
)
)
layer.register_parameter("w13_weight", w13_weight)
Expand All @@ -647,6 +660,7 @@ def create_weights(
hidden_size,
intermediate_size_per_partition,
dtype=params_dtype,
device=weight_device,
)
)
layer.register_parameter("w2_weight", w2_weight)
Expand Down Expand Up @@ -911,26 +925,49 @@ def create_weights(
self.intermediate_pad = (
self.intermediate_size - layer.intermediate_size_per_partition
)
# Streaming source buffers are transient: the loader fills them from the
# (logical, unpadded) MXFP4 checkpoint, _online_quant reads them once,
# then they're freed. They never feed a kernel, so they skip the
# kernel-alignment padding the real (target) buffers get, and are
# meta-allocated (0 real memory at build time). Using the logical size
# keeps the loader's copy-count == expected numel so streaming actually
# triggers -- otherwise the padded shapes would never fill and the layer
# would silently fall back to the two-pass path. _online_quant clears
# _stream_online before re-creating the target buffers, so those stay
# real and padded. This readies the MXFP4-source streaming path for a
# future online target (e.g. nvfp4).
stream_src = getattr(layer, "_stream_online", False)
weight_device = "meta" if stream_src else None
int_dim = (
intermediate_size_per_partition
if stream_src
else intermediate_size_per_partition_after_pad
)
hid_dim = layer.hidden_size if stream_src else hidden_size
# Fused gate_up_proj (column parallel)
w13_weight = atom_parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition_after_pad, # TP included
hidden_size // 2,
2 * int_dim, # TP included
hid_dim // 2,
dtype=weight_dtype,
device=weight_device,
)
)
layer.register_parameter("w13_weight", w13_weight)
# Zero-fill padding region: FP4 dtype doesn't support torch.zeros,
# so we zero the underlying bytes to avoid garbage in padded rows.
w13_weight.data.view(torch.uint8).zero_()
# Zero-fill padding region: FP4 dtype doesn't support torch.zeros, so we
# zero the underlying bytes to avoid garbage in padded rows. Skipped for
# streaming sources: they are meta (can't be zeroed here; the loader
# zeroes on materialization) and unpadded (nothing to zero).
if not stream_src:
w13_weight.data.view(torch.uint8).zero_()
set_weight_attrs(w13_weight, extra_weight_attrs)

w13_weight_scale = atom_parameter(
torch.zeros(
num_experts,
2 * intermediate_size_per_partition_after_pad,
hidden_size // mxfp4_block,
2 * int_dim,
hid_dim // mxfp4_block,
dtype=scale_dtype,
)
)
Expand All @@ -941,7 +978,7 @@ def create_weights(
w13_bias = atom_parameter(
torch.zeros(
num_experts,
2 * intermediate_size_per_partition_after_pad,
2 * int_dim,
dtype=torch.bfloat16,
)
)
Expand All @@ -954,20 +991,22 @@ def create_weights(
w2_weight = atom_parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition_after_pad // 2, # TP included
hid_dim,
int_dim // 2, # TP included
dtype=weight_dtype,
device=weight_device,
)
)
layer.register_parameter("w2_weight", w2_weight)
w2_weight.data.view(torch.uint8).zero_()
if not stream_src:
w2_weight.data.view(torch.uint8).zero_()
set_weight_attrs(w2_weight, extra_weight_attrs)

w2_weight_scale = atom_parameter(
torch.zeros(
num_experts,
hidden_size,
intermediate_size_per_partition_after_pad // mxfp4_block,
hid_dim,
int_dim // mxfp4_block,
dtype=scale_dtype,
)
)
Expand All @@ -978,7 +1017,7 @@ def create_weights(
w2_bias = atom_parameter(
torch.zeros(
num_experts,
hidden_size,
hid_dim,
dtype=torch.bfloat16,
)
)
Expand Down Expand Up @@ -1940,6 +1979,15 @@ def create_weights(
):
self.num_experts = num_experts
intermediate_size_for_weight = intermediate_size_per_partition
# Streaming source buffers are transient: the loader fills them from the
# checkpoint, _online_quant reads them once, then they're freed. They
# never feed a kernel, so they don't need the kernel-alignment padding
# the real (target) buffers get. Allocating them at the logical
# (checkpoint) size keeps the loader's copy-count == expected numel so
# streaming actually triggers -- otherwise a padded per_1x32 source
# (e.g. MXFP8 with TP-misaligned intermediate) would never reach the
# padded numel and silently fall back to the two-pass path.
stream_src = getattr(layer, "_stream_online", False)

if self.block_quant:
if self.quant_type == QuantType.per_1x128:
Expand All @@ -1966,10 +2014,13 @@ def create_weights(
f"{intermediate_size_per_partition} is not divisible by "
f"weight quantization block_k = {block_k}."
)
if self.quant_type == QuantType.per_1x32:
if self.quant_type == QuantType.per_1x32 and not stream_src:
# aiter's GU-interleaved MXFP8 scale shuffle packs 8 scale
# columns, i.e. 256 weight columns for 1x32 scales. TP8 on
# MiniMax-M3 has local intermediate=384, so pad to 512.
# Skipped for streaming source buffers (see stream_src above):
# they only hold the checkpoint payload and are freed after
# _online_quant, so they stay at the logical size.
scale_pack_k = block_k * 8
intermediate_size_for_weight = (
(intermediate_size_per_partition + scale_pack_k - 1)
Expand All @@ -1978,16 +2029,27 @@ def create_weights(
)

# WEIGHTS
# When this module streams its online re-quant (e.g. ptpc_fp8 -> mxfp4,
# MXFP8 -> mxfp4), allocate the source w13/w2 buffers on the meta device
# (0 real memory at build time); the loader materializes them on first
# touch and _online_quant frees them right after re-quantizing. Only the
# initial (source) creation is meta -- _online_quant clears the flag
# before re-invoking create_weights for the FP8 target so those stay
# real (and get the kernel-alignment padding skipped above).
weight_device = "meta" if stream_src else None
w13_weight = atom_parameter(
torch.empty(
num_experts,
2 * intermediate_size_for_weight,
hidden_size,
dtype=params_dtype,
device=weight_device,
)
)
layer.register_parameter("w13_weight", w13_weight)
if self.quant_type == QuantType.per_1x32:
# Zero padding bytes for the real (non-streaming) padded per_1x32 buffer;
# streaming source buffers are unpadded (stream_src) so nothing to zero.
if self.quant_type == QuantType.per_1x32 and not stream_src:
w13_weight.data.view(torch.uint8).zero_()
set_weight_attrs(w13_weight, extra_weight_attrs)

Expand All @@ -1997,10 +2059,11 @@ def create_weights(
hidden_size,
intermediate_size_for_weight,
dtype=params_dtype,
device=weight_device,
)
)
layer.register_parameter("w2_weight", w2_weight)
if self.quant_type == QuantType.per_1x32:
if self.quant_type == QuantType.per_1x32 and not stream_src:
w2_weight.data.view(torch.uint8).zero_()
set_weight_attrs(w2_weight, extra_weight_attrs)

Expand Down Expand Up @@ -2658,6 +2721,22 @@ def __init__(
"params_dtype": self.params_dtype,
"weight_loader": self.weight_loader,
}
# Stream online quantization: quantize this expert module right after its
# weights finish loading (freeing the source BF16), same decision/flow as
# LinearBase. Must be set BEFORE create_weights so the unquantized method
# allocates the source w13/w2 buffers on the meta device; the loader
# materializes them on first touch (see model_loader/loader.py). The
# target buffers allocated later inside _online_quant stay real.
self._stream_online = should_stream_online_quant(
quant_config,
prefix,
layer_quant_config.quant_type if layer_quant_config else QuantType.No,
self.params_dtype,
# FusedMoE._online_quant only dequantizes No/per_Token/per_1x128/
# per_1x32 sources -- per_Tensor is Linear-only.
allow_per_tensor_source=False,
)
self._load_device = torch.empty(0).device if self._stream_online else None
self.quant_method.create_weights(layer=self, **self.moe_quant_params)
compilation_config = atom_config.compilation_config
if prefix in compilation_config.static_forward_context:
Expand Down Expand Up @@ -2772,6 +2851,10 @@ def check_need_allgather():
f"Unsupported online quant_dtype for MoE: {online_quant_dtype}"
)
self.moe_quant_params["params_dtype"] = online_quant_dtype
# The source buffers were meta-allocated for streaming; the target
# buffers we build now must be real, so clear the flag before the second
# create_weights (Fp8MoEMethod reads it to decide meta allocation).
self._stream_online = False
with torch.device(device):
self.quant_method.create_weights(layer=self, **self.moe_quant_params)

Expand Down
Loading