feat: [Qwen3.5] Support block-wise FP8 quantization and model adaptation (#18926)
This commit is contained in:
@@ -728,6 +728,51 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
|
|||||||
)
|
)
|
||||||
self.weight_loader_v2(param, loaded_weight_shard, shard_id)
|
self.weight_loader_v2(param, loaded_weight_shard, shard_id)
|
||||||
|
|
||||||
|
def _load_merged_block_scale(
|
||||||
|
self, param: BasevLLMParameter, loaded_weight: torch.Tensor
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Handle block-wise scale loading for MergedColumnParallelLinear.
|
||||||
|
Similar to QKVParallelLinear._load_qkv_block_scale, but for merged column layers.
|
||||||
|
"""
|
||||||
|
weight_block_size = self.quant_method.quant_config.weight_block_size
|
||||||
|
block_n, _ = weight_block_size[0], weight_block_size[1]
|
||||||
|
block_n = 1 if getattr(param, "format_ue8m0", False) else block_n
|
||||||
|
|
||||||
|
# Calculate block sizes for each shard
|
||||||
|
shard_block_sizes = []
|
||||||
|
shard_block_offsets = []
|
||||||
|
current_block_offset = 0
|
||||||
|
for output_size in self.output_sizes:
|
||||||
|
shard_block_size = (output_size + block_n - 1) // block_n
|
||||||
|
shard_block_sizes.append(shard_block_size)
|
||||||
|
shard_block_offsets.append(current_block_offset)
|
||||||
|
current_block_offset += shard_block_size
|
||||||
|
|
||||||
|
# Load each shard
|
||||||
|
for shard_id, (shard_block_offset, shard_block_size) in enumerate(
|
||||||
|
zip(shard_block_offsets, shard_block_sizes)
|
||||||
|
):
|
||||||
|
# Extract the shard from loaded_weight
|
||||||
|
loaded_weight_shard = loaded_weight.narrow(
|
||||||
|
param.output_dim, shard_block_offset, shard_block_size
|
||||||
|
)
|
||||||
|
|
||||||
|
# Calculate per-rank offset and size (considering TP)
|
||||||
|
rank_shard_offset = shard_block_offset // self.tp_size
|
||||||
|
rank_shard_size = shard_block_size // self.tp_size
|
||||||
|
|
||||||
|
# Load into the parameter
|
||||||
|
param.load_merged_column_weight(
|
||||||
|
loaded_weight=loaded_weight_shard,
|
||||||
|
shard_id=shard_id,
|
||||||
|
shard_offset=rank_shard_offset,
|
||||||
|
shard_size=rank_shard_size,
|
||||||
|
tp_rank=self.tp_rank,
|
||||||
|
tp_size=self.tp_size,
|
||||||
|
use_presharded_weights=self.use_presharded_weights,
|
||||||
|
)
|
||||||
|
|
||||||
def weight_loader_v2(
|
def weight_loader_v2(
|
||||||
self,
|
self,
|
||||||
param: BasevLLMParameter,
|
param: BasevLLMParameter,
|
||||||
@@ -743,6 +788,9 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
|
|||||||
tp_size=self.tp_size,
|
tp_size=self.tp_size,
|
||||||
)
|
)
|
||||||
return
|
return
|
||||||
|
elif isinstance(param, BlockQuantScaleParameter):
|
||||||
|
self._load_merged_block_scale(param, loaded_weight)
|
||||||
|
return
|
||||||
elif type(param) in (RowvLLMParameter, BasevLLMParameter):
|
elif type(param) in (RowvLLMParameter, BasevLLMParameter):
|
||||||
param.load_merged_column_weight(
|
param.load_merged_column_weight(
|
||||||
loaded_weight=loaded_weight,
|
loaded_weight=loaded_weight,
|
||||||
|
|||||||
@@ -171,8 +171,11 @@ class Fp8Config(QuantizationConfig):
|
|||||||
config, ["ignored_layers", "modules_to_not_convert"], None
|
config, ["ignored_layers", "modules_to_not_convert"], None
|
||||||
)
|
)
|
||||||
if ignored_layers:
|
if ignored_layers:
|
||||||
# hack for ministral
|
if "mistral3" in config.get("model_type", ""):
|
||||||
ignored_layers = [layer.replace("model.", "") for layer in ignored_layers]
|
# hack for ministral
|
||||||
|
ignored_layers = [
|
||||||
|
layer.replace("model.", "") for layer in ignored_layers
|
||||||
|
]
|
||||||
weight_block_size = cls.get_from_keys_or(config, ["weight_block_size"], None)
|
weight_block_size = cls.get_from_keys_or(config, ["weight_block_size"], None)
|
||||||
if use_mxfp8 and weight_block_size is not None:
|
if use_mxfp8 and weight_block_size is not None:
|
||||||
logger.warning(
|
logger.warning(
|
||||||
|
|||||||
@@ -64,7 +64,7 @@ class Qwen3_5ForCausalLMMTP(nn.Module):
|
|||||||
self.model = Qwen3_5ForCausalLM(
|
self.model = Qwen3_5ForCausalLM(
|
||||||
config,
|
config,
|
||||||
quant_config,
|
quant_config,
|
||||||
prefix=add_prefix("model", prefix),
|
prefix=add_prefix("mtp", prefix),
|
||||||
)
|
)
|
||||||
|
|
||||||
if get_pp_group().is_last_rank:
|
if get_pp_group().is_last_rank:
|
||||||
@@ -214,16 +214,11 @@ class Qwen3_5ForCausalLMMTP(nn.Module):
|
|||||||
if "mtp" not in name:
|
if "mtp" not in name:
|
||||||
continue
|
continue
|
||||||
|
|
||||||
# Some checkpoints use model.language_model.mtp.* prefix
|
|
||||||
if "language_model" in name:
|
|
||||||
name = name.replace(r"model.language_model.", r"model.")
|
|
||||||
|
|
||||||
if name.startswith("mtp."):
|
if name.startswith("mtp."):
|
||||||
# Remove the mtp. prefix for processing
|
# Remove the mtp. prefix for processing
|
||||||
name = name.replace("mtp.", "model.")
|
name = name.replace("mtp.", "model.")
|
||||||
|
|
||||||
name = name.replace("model.fc", "fc")
|
name = name.replace("model.fc", "fc")
|
||||||
name = name.replace("model.norm", "norm")
|
|
||||||
name = name.replace("model.pre_fc", "pre_fc")
|
name = name.replace("model.pre_fc", "pre_fc")
|
||||||
|
|
||||||
if ".self_attn." in name:
|
if ".self_attn." in name:
|
||||||
|
|||||||
@@ -786,9 +786,9 @@ class Qwen3VLForConditionalGeneration(nn.Module):
|
|||||||
config.vision_config,
|
config.vision_config,
|
||||||
# NOTE: Qwen3-VL vision encoder currently supports BitsAndBytes 4-bit quantization.
|
# NOTE: Qwen3-VL vision encoder currently supports BitsAndBytes 4-bit quantization.
|
||||||
# Other quantization methods (e.g., GPTQ, AWQ) are untested and may not be supported.
|
# Other quantization methods (e.g., GPTQ, AWQ) are untested and may not be supported.
|
||||||
quant_config=quant_config,
|
quant_config=None,
|
||||||
norm_eps=getattr(config, "rms_norm_eps", 1e-6),
|
norm_eps=getattr(config, "rms_norm_eps", 1e-6),
|
||||||
prefix=add_prefix("visual", prefix),
|
prefix=add_prefix("model.visual", prefix),
|
||||||
use_data_parallel=self.use_data_parallel,
|
use_data_parallel=self.use_data_parallel,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -804,7 +804,7 @@ class Qwen3VLForConditionalGeneration(nn.Module):
|
|||||||
self.model = language_model_cls(
|
self.model = language_model_cls(
|
||||||
config=self.config,
|
config=self.config,
|
||||||
quant_config=quant_config,
|
quant_config=quant_config,
|
||||||
prefix=add_prefix("model", prefix),
|
prefix=add_prefix("model.language_model", prefix),
|
||||||
)
|
)
|
||||||
if self.pp_group.is_last_rank:
|
if self.pp_group.is_last_rank:
|
||||||
if self.pp_group.world_size == 1 and self.config.tie_word_embeddings:
|
if self.pp_group.world_size == 1 and self.config.tie_word_embeddings:
|
||||||
@@ -1110,7 +1110,6 @@ class Qwen3VLForConditionalGeneration(nn.Module):
|
|||||||
if "visual" in name:
|
if "visual" in name:
|
||||||
# adapt to VisionAttention
|
# adapt to VisionAttention
|
||||||
name = name.replace(r"attn.qkv.", r"attn.qkv_proj.")
|
name = name.replace(r"attn.qkv.", r"attn.qkv_proj.")
|
||||||
name = name.replace(r"model.visual.", r"visual.")
|
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# Skip loading extra bias for GPTQ models.
|
# Skip loading extra bias for GPTQ models.
|
||||||
|
|||||||
Reference in New Issue
Block a user