319 lines
15 KiB
Python
319 lines
15 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
from .fp8_optimization import _is_fsdp_managed
|
|
|
|
# E2M1: 1 sign bit, 2 exponent bits, 1 mantissa bit -> 8 magnitudes, 16 codes.
|
|
FLOAT4_MAX = 6.0
|
|
# NVFP4 blocks are 16 elements wide and scale them with a single E4M3 byte (plus one
|
|
# fp32 global scale per tensor); block 32 with an fp32 scale is the OCP MXFP4 flavour.
|
|
FLOAT4_BLOCK_SIZE = 16
|
|
FLOAT8_E4M3 = torch.float8_e4m3fn
|
|
FLOAT8_E4M3_MAX = torch.finfo(FLOAT8_E4M3).max # 448.0
|
|
FLOAT4_MAGNITUDES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
|
|
# Mid-points between consecutive magnitudes; a value is rounded to the bucket this
|
|
# falls into. Boundaries are exclusive for the lower bucket (see _float4_round).
|
|
FLOAT4_BOUNDARIES = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0)
|
|
# Index = code (magnitude index | sign << 3).
|
|
FLOAT4_LUT = FLOAT4_MAGNITUDES + tuple(-v for v in FLOAT4_MAGNITUDES)
|
|
|
|
_FLOAT4_TABLES = {}
|
|
|
|
|
|
def _float4_tables(device):
|
|
"""The two constant tables, built once per device."""
|
|
tables = _FLOAT4_TABLES.get(device)
|
|
if tables is None:
|
|
tables = (
|
|
torch.tensor(FLOAT4_LUT, dtype=torch.float32, device=device),
|
|
torch.tensor(FLOAT4_MAGNITUDES, dtype=torch.float32, device=device),
|
|
torch.tensor(FLOAT4_BOUNDARIES, dtype=torch.float32, device=device),
|
|
)
|
|
_FLOAT4_TABLES[device] = tables
|
|
return tables
|
|
|
|
|
|
def _float4_scale_name(param_name):
|
|
return param_name + "_fp4_scale"
|
|
|
|
|
|
def _float4_global_scale_name(param_name):
|
|
# NVFP4's second-level, per-tensor scale (a single fp32 value), stored next to the
|
|
# per-block E4M3 scales so both move together under the offload hooks.
|
|
return param_name + "_fp4_global_scale"
|
|
|
|
|
|
def _float4_meta_name(param_name):
|
|
# Plain attribute (not a buffer): the unpacked shape and the block size are Python
|
|
# constants that never need to move between devices.
|
|
return param_name + "_fp4_meta"
|
|
|
|
|
|
def _float4_round(magnitudes, boundaries):
|
|
"""Round non-negative scaled values to the nearest E2M1 magnitude index.
|
|
|
|
``bucketize(..., right=True)`` is round-half-up; a value sitting exactly on a
|
|
boundary is then pushed down to the even code, which makes the whole thing
|
|
round-nearest-even like the hardware FP8 cast in fp8_optimization. Adjacent codes
|
|
always differ in parity, so on a tie either neighbour can be the even one and one
|
|
where covers both directions.
|
|
"""
|
|
idx = torch.bucketize(magnitudes, boundaries, right=True)
|
|
lower = (idx - 1).clamp(min=0)
|
|
# `lower` indexes the boundary between the two candidate codes, not a grid value.
|
|
tie = (magnitudes == boundaries[lower]) & (idx > 0)
|
|
return torch.where(tie & (lower % 2 == 0), lower, idx)
|
|
|
|
|
|
@torch.no_grad()
|
|
def quantize_weight_to_float4(tensor, block_size=FLOAT4_BLOCK_SIZE):
|
|
"""Return ``(packed_uint8, block_scale_fp8, global_scale)`` for a 2-D weight.
|
|
|
|
NVFP4's two-level scaling: the per-tensor ``global_scale`` maps the largest block
|
|
scale onto the E4M3 grid, and each 16-element block stores its scale as one E4M3
|
|
byte. The element codes are quantized with the *dequantized* block scale
|
|
(``block_scale_fp8.float() / global_scale``), so the stored nibbles reproduce
|
|
exactly what the unpack divides back out -- no scale mismatch between the two.
|
|
"""
|
|
rows, features = tensor.shape[0], tensor.shape[1]
|
|
x = tensor.detach().float().reshape(rows, features)
|
|
pad = (-features) % block_size
|
|
if pad:
|
|
# Only reachable for weights whose input dim is not a multiple of the block;
|
|
# the padding is dropped again by the unpack through the stored shape.
|
|
x = nn.functional.pad(x, (0, pad))
|
|
padded = x.shape[1]
|
|
blocks = x.reshape(-1, block_size)
|
|
|
|
# Second-level (per-tensor) scale: (E4M3_MAX * E2M1_MAX) / global_amax, so the
|
|
# widest block (amax == global_amax) scales to exactly E4M3_MAX and no block scale
|
|
# overflows fp8. global_amax is clamped so an all-zero weight yields codes of 0.
|
|
global_amax = x.abs().amax().clamp(min=1e-12)
|
|
global_scale = (FLOAT8_E4M3_MAX * FLOAT4_MAX) / global_amax
|
|
|
|
# First-level (per-block) scale, rounded to E4M3 (NVFP4's storage type) then read
|
|
# back, so quantization uses the identical scale the unpack will use. A block whose
|
|
# scale underflows to 0 in fp8 (or a zero-initialized control projection) is clamped
|
|
# instead of dividing by zero, keeping its codes at 0 rather than NaN.
|
|
block_amax = blocks.abs().amax(dim=1, keepdim=True)
|
|
block_scale_fp8 = ((block_amax / FLOAT4_MAX) * global_scale).to(FLOAT8_E4M3)
|
|
block_scale = (block_scale_fp8.float() / global_scale).clamp(min=1e-12)
|
|
|
|
_, _, boundaries = _float4_tables(x.device)
|
|
scaled = (blocks / block_scale).clamp(-FLOAT4_MAX, FLOAT4_MAX)
|
|
idx = _float4_round(scaled.abs().reshape(rows, padded), boundaries)
|
|
code = torch.where(scaled.reshape(rows, padded) < 0, idx + 8, idx).to(torch.uint8)
|
|
packed = code[:, 0::2] | (code[:, 1::2] << 4)
|
|
return packed, block_scale_fp8.reshape(rows, padded // block_size, 1), global_scale
|
|
|
|
|
|
@torch.no_grad()
|
|
def dequantize_weight_from_float4(packed, block_scale_fp8, global_scale, shape,
|
|
block_size=FLOAT4_BLOCK_SIZE, dtype=torch.bfloat16):
|
|
"""Rebuild the ``(out, in)`` weight from the nibbles and the two-level scales."""
|
|
rows, half = packed.shape
|
|
code = torch.stack((packed & 0x0F, packed >> 4), dim=-1).reshape(rows, half * 2)
|
|
lut, _, _ = _float4_tables(packed.device)
|
|
values = lut[code.long()]
|
|
# E4M3 block scale back to fp32, then undo the per-tensor global scale.
|
|
block_scale = block_scale_fp8.float() / global_scale
|
|
blocks = values.reshape(rows, block_scale_fp8.shape[1], block_size) * block_scale
|
|
return blocks.reshape(rows, -1)[:, :shape[1]].reshape(shape).to(dtype)
|
|
|
|
|
|
def _float4_pack_weight(module, device, block_size):
|
|
"""Pack ``module.weight`` (a 2-D floating tensor) into E2M1 nibbles + scale/meta buffers.
|
|
|
|
Single source of truth for the packing, shared by the whole-model
|
|
:func:`convert_model_weight_to_float4` and the per-layer
|
|
:func:`float4_quantize_linear` used around a LoRA merge.
|
|
"""
|
|
weight = module.weight
|
|
packed, block_scale_fp8, global_scale = quantize_weight_to_float4(weight.data, block_size)
|
|
to_device = (lambda t: t.to(device)) if device is not None else (lambda t: t)
|
|
module.register_buffer(_float4_scale_name("weight"), to_device(block_scale_fp8))
|
|
module.register_buffer(_float4_global_scale_name("weight"), to_device(global_scale))
|
|
setattr(module, _float4_meta_name("weight"), (tuple(weight.shape), block_size))
|
|
# A Parameter may only hold a non-floating tensor while it does not require
|
|
# grad, and a nibble-packed weight could not be trained through anyway.
|
|
weight.requires_grad_(False)
|
|
weight.data = to_device(packed)
|
|
|
|
|
|
def _float4_wrap_forward(module, origin_dtype):
|
|
"""Install the per-forward FP4 unpack on ONE module holding a packed weight."""
|
|
if getattr(module, _float4_meta_name("weight"), None) is None:
|
|
return
|
|
if hasattr(module, "original_forward_fp4"):
|
|
return
|
|
module.original_forward_fp4 = module.forward
|
|
setattr(
|
|
module,
|
|
"forward",
|
|
lambda *inputs, m=module, **kwargs: _float4_autocast_forward(m, origin_dtype, *inputs, **kwargs)
|
|
)
|
|
|
|
|
|
def convert_model_weight_to_float4(model, exclude_module_name=("modulation",),
|
|
device=None, block_size=FLOAT4_BLOCK_SIZE):
|
|
"""Pack every ``nn.Linear`` weight of ``model`` into E2M1 nibbles, in place."""
|
|
if _is_fsdp_managed(model):
|
|
raise NotImplementedError(
|
|
"GPU_memory_mode 'qfloat4' is not supported under FSDP: the managed "
|
|
"parameters are flat-storage views, so a block-scaled 4-bit repacking of "
|
|
"`weight` has no shape to keep. Use 'qfloat8' (which dequantizes on the "
|
|
"output side instead) or no quantization.")
|
|
for name, module in model.named_modules():
|
|
if not isinstance(module, nn.Linear):
|
|
continue
|
|
if any(keyword in name for keyword in exclude_module_name):
|
|
continue
|
|
weight = module.weight
|
|
if weight is None or weight.dtype == torch.uint8 or weight.dim() != 2:
|
|
continue
|
|
_float4_pack_weight(module, device, block_size)
|
|
return model
|
|
|
|
|
|
@torch.no_grad()
|
|
def _float4_autocast_forward(module, origin_dtype, *inputs, **kwargs):
|
|
"""Swap the packed weight for its bf16 reconstruction around one forward.
|
|
|
|
Unlike the FP8 wrapper this never re-quantizes on the way out: the packed nibbles
|
|
are simply put back, so the stored bits cannot drift across steps and the dequant
|
|
costs one gather plus one mul instead of a division.
|
|
"""
|
|
restored = []
|
|
for param_name, param in list(module.named_parameters(recurse=False)):
|
|
meta = getattr(module, _float4_meta_name(param_name), None)
|
|
if meta is None:
|
|
if param.is_floating_point() and param.dtype != origin_dtype:
|
|
# An unquantized fp32 bias next to a bf16 weight would make F.linear
|
|
# reject the pair; fp8_optimization normalizes them the same way.
|
|
restored.append((param, param.data))
|
|
param.data = param.data.to(origin_dtype)
|
|
continue
|
|
if param.dtype != torch.uint8:
|
|
# `module.to(dtype)` walks every parameter and would silently turn the
|
|
# packed bytes into the bf16 numbers 0..255. Refuse instead of generating
|
|
# with a destroyed weight.
|
|
raise RuntimeError(
|
|
"The FP4-packed parameter '%s' holds a %s tensor, not the packed uint8 "
|
|
"nibbles: something called .to(dtype) on a qfloat4 model. Apply the "
|
|
"memory mode after the dtype move, and move devices with "
|
|
"apply_gpu_memory_mode()/offload hooks, which keep the dtype."
|
|
% (param_name, param.dtype))
|
|
shape, block_size = meta
|
|
restored.append((param, param.data))
|
|
param.data = dequantize_weight_from_float4(
|
|
param.data, getattr(module, _float4_scale_name(param_name)),
|
|
getattr(module, _float4_global_scale_name(param_name)),
|
|
shape, block_size, origin_dtype)
|
|
|
|
# Integer tensors (indices) must keep their dtype or the kernel rejects them.
|
|
inputs = [input.to(origin_dtype) if torch.is_tensor(input) and input.is_floating_point()
|
|
else input for input in inputs]
|
|
kwargs = {key: value.to(origin_dtype) if torch.is_tensor(value) and value.is_floating_point()
|
|
else value for key, value in kwargs.items()}
|
|
try:
|
|
return module.original_forward_fp4(*inputs, **kwargs)
|
|
finally:
|
|
for param, data in restored:
|
|
param.data = data
|
|
|
|
|
|
def convert_float4_weight_dtype_wrapper(module, origin_dtype):
|
|
"""Install the per-forward unpack on every module holding a packed weight."""
|
|
for name, child in module.named_modules():
|
|
if name == "" or "embed_tokens" in name or isinstance(child, nn.Embedding):
|
|
continue
|
|
_float4_wrap_forward(child, origin_dtype)
|
|
return module
|
|
|
|
|
|
def undo_convert_float4_weight_dtype_wrapper(module, origin_dtype):
|
|
"""Drop the wrappers and unpack in place, leaving a plain full-precision module.
|
|
|
|
The mirror of the FP8 undo. Afterwards the model runs its own forwards again and no
|
|
per-forward reconstruction is needed -- at the cost of holding the weights in
|
|
`origin_dtype` again, so this is what to call before a LoRA merge or a save.
|
|
"""
|
|
for _, child in module.named_modules():
|
|
if not hasattr(child, "original_forward_fp4"):
|
|
continue
|
|
setattr(child, "forward", child.original_forward_fp4)
|
|
delattr(child, "original_forward_fp4")
|
|
meta = getattr(child, _float4_meta_name("weight"), None)
|
|
if meta is None:
|
|
continue
|
|
shape, block_size = meta
|
|
weight = child.weight
|
|
scale = getattr(child, _float4_scale_name("weight"), None)
|
|
global_scale = getattr(child, _float4_global_scale_name("weight"), None)
|
|
if weight is not None and weight.dtype == torch.uint8 and scale is not None:
|
|
weight.requires_grad_(False)
|
|
weight.data = dequantize_weight_from_float4(
|
|
weight.data, scale, global_scale, shape, block_size, origin_dtype)
|
|
delattr(child, _float4_meta_name("weight"))
|
|
if scale is not None:
|
|
delattr(child, _float4_scale_name("weight"))
|
|
if global_scale is not None:
|
|
delattr(child, _float4_global_scale_name("weight"))
|
|
return module
|
|
|
|
|
|
def float4_dequantize_linear(module, origin_dtype):
|
|
"""Unpack ONE FP4-packed ``nn.Linear`` back to ``origin_dtype``, in place.
|
|
|
|
Drops this layer's scale/meta buffers and forward wrapper, leaving a plain
|
|
full-precision Linear. Returns True if the layer was packed (so the caller knows to
|
|
re-pack it), False if it was already full precision and was left untouched.
|
|
|
|
Deliberately per layer rather than whole-model: a LoRA merge edits one layer at a
|
|
time, and materializing the *entire* transformer in ``origin_dtype`` would need the
|
|
very memory the 4-bit packing was chosen to avoid. Only this layer's weight is
|
|
dequantized, so the peak overhead is one layer.
|
|
"""
|
|
meta = getattr(module, _float4_meta_name("weight"), None)
|
|
if meta is None:
|
|
return False
|
|
if _is_fsdp_managed(module):
|
|
raise NotImplementedError(
|
|
"Cannot dequantize an FP4-packed weight under FSDP: the managed parameters "
|
|
"are flat-storage DTensor views with no shape to unpack. Merge the LoRA "
|
|
"before the FSDP sharding (or before apply_gpu_memory_mode).")
|
|
shape, block_size = meta
|
|
weight = module.weight
|
|
scale = getattr(module, _float4_scale_name("weight"), None)
|
|
global_scale = getattr(module, _float4_global_scale_name("weight"), None)
|
|
if hasattr(module, "original_forward_fp4"):
|
|
setattr(module, "forward", module.original_forward_fp4)
|
|
delattr(module, "original_forward_fp4")
|
|
if weight is not None and weight.dtype == torch.uint8 and scale is not None:
|
|
weight.requires_grad_(False)
|
|
weight.data = dequantize_weight_from_float4(
|
|
weight.data, scale, global_scale, shape, block_size, origin_dtype)
|
|
delattr(module, _float4_meta_name("weight"))
|
|
if scale is not None:
|
|
delattr(module, _float4_scale_name("weight"))
|
|
if global_scale is not None:
|
|
delattr(module, _float4_global_scale_name("weight"))
|
|
return True
|
|
|
|
|
|
def float4_quantize_linear(module, origin_dtype, device=None, block_size=FLOAT4_BLOCK_SIZE):
|
|
"""Re-pack ONE ``nn.Linear`` into FP4 and reinstall its forward wrapper.
|
|
|
|
The exact inverse of :func:`float4_dequantize_linear`, called after a LoRA delta has
|
|
been folded into the layer's full-precision weight so the model stays 4-bit for
|
|
generation. Returns True if the layer was packed.
|
|
"""
|
|
if not isinstance(module, nn.Linear):
|
|
return False
|
|
weight = module.weight
|
|
if weight is None or weight.dtype == torch.uint8 or weight.dim() != 2:
|
|
return False
|
|
_float4_pack_weight(module, device, block_size)
|
|
_float4_wrap_forward(module, origin_dtype)
|
|
return True
|