diff --git a/examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yaml b/examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yaml index 68d7e4f1fa19..e8fc2f04f468 100644 --- a/examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yaml +++ b/examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yaml @@ -29,10 +29,16 @@ transforms: enabled: true fuse_nvfp4_moe: backend: trtllm_gen + detect_sharding: + enabled: false + sharding_transform_executor: + enabled: false apply_sharding_hints: + enabled: true allreduce_strategy: SYMM_MEM # Shared expert is excluded from sharding for performance purpose shard_layers: ["moe", "delta", "mha"] + simple_shard_filter: "lm_head" multi_stream_moe: stage: compile enabled: true diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py index 7c6071081fa1..76e0a6483cff 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py @@ -635,26 +635,25 @@ def __init__(self, config: Qwen3_5MoeTextConfig, intermediate_size: int): self.act_fn = ACT2FN[config.hidden_act] def forward(self, x: torch.Tensor) -> torch.Tensor: + # Intentionally left untagged: with no ``layer_type`` it defaults to "unknown", + # which the ``shard_layers`` inclusion whitelist excludes -> kept replicated. gate = torch.ops.auto_deploy.torch_linear_simple( x, self.gate_proj.weight, self.gate_proj.bias, tp_mode="colwise", - layer_type="shared_expert", ) up = torch.ops.auto_deploy.torch_linear_simple( x, self.up_proj.weight, self.up_proj.bias, tp_mode="colwise", - layer_type="shared_expert", ) return torch.ops.auto_deploy.torch_linear_simple( self.act_fn(gate) * up, self.down_proj.weight, self.down_proj.bias, tp_mode="rowwise", - layer_type="shared_expert", ) diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_swiglu.py b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_swiglu.py index 781910fa8928..700cb3693b27 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_swiglu.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_swiglu.py @@ -22,7 +22,6 @@ The SwiGLU pattern is: silu(x @ gate.T) * (x @ up.T) @ down.T """ -from contextlib import contextmanager from typing import Tuple, Type import torch @@ -48,7 +47,7 @@ eliminate_dead_code, get_attr_by_name, ) -from ...utils.node_utils import extract_op_args, is_op, set_op_args +from ...utils.node_utils import is_op from ...utils.pattern_matcher import ADPatternMatcherPass, register_ad_pattern from ...utils.quantization_utils import ensure_tma_col_major from ..interface import ( @@ -60,42 +59,6 @@ ) -def _weight_key(node: Node): - """Stable key for a weight arg: the get_attr target FQN (survives node re-creation - during pattern replacement), falling back to node identity. ``args[1]`` is the - weight for both linear and SwiGLU ops.""" - w = node.args[1] if len(node.args) > 1 else None - if not isinstance(w, Node): - return None - return w.target if w.op == "get_attr" else w - - -@contextmanager -def preserve_layer_types(gm: GraphModule, linear_op, fused_op): - """Carry the ``layer_type`` hint across a fusion that consumes ``linear_op`` nodes - and emits ``fused_op`` nodes (which would otherwise drop the hint). - - Snapshots each source weight's ``layer_type`` before the rewrite, then re-applies - it to the fused node keyed by weight, so hint-driven sharding (``shard_layers``) - can still classify the fused node. Wrap the matcher's ``patterns.apply`` call. - """ - wmap = {} - for n in gm.graph.nodes: - if is_op(n, linear_op): - [lt] = extract_op_args(n, "layer_type") - key = _weight_key(n) - if lt is not None and key is not None: - wmap[key] = lt - yield - if not wmap: - return - for n in gm.graph.nodes: - if is_op(n, fused_op): - key = _weight_key(n) - if key is not None and key in wmap: - set_op_args(n, layer_type=wmap[key]) - - def _maybe_to_deepgemm_layout( weight: torch.Tensor, scale: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: @@ -279,12 +242,7 @@ def _apply( dummy_args=dummy_args_with_bias, ) - with preserve_layer_types( - gm, - torch.ops.auto_deploy.torch_linear_simple.default, - torch.ops.auto_deploy.torch_swiglu_mlp.default, - ): - num_matches = patterns.apply(gm.graph) + num_matches = patterns.apply(gm.graph) if num_matches > 0: gm.recompile() @@ -596,12 +554,7 @@ def _apply( dummy_args=dummy_args, ) - with preserve_layer_types( - gm, - torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear.default, - torch.ops.auto_deploy.torch_nvfp4_swiglu_mlp.default, - ): - num_matches = patterns.apply(gm.graph) + num_matches = patterns.apply(gm.graph) if num_matches > 0: gm.recompile() @@ -882,12 +835,7 @@ def _apply( dummy_args=dummy_args, ) - with preserve_layer_types( - gm, - torch.ops.auto_deploy.torch_fake_quant_finegrained_fp8_linear.default, - torch.ops.auto_deploy.torch_finegrained_fp8_swiglu_mlp.default, - ): - num_matches = patterns.apply(gm.graph) + num_matches = patterns.apply(gm.graph) if num_matches > 0: gm.recompile() diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding_ir.py b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding_ir.py index f191a817c13e..88b1aef1ce00 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding_ir.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding_ir.py @@ -148,6 +148,37 @@ def _fp4_weight_scale_pipeline_cache_spec( } +def _rowwise_bias_load_hook(state_dict, prefix, *args, rank, param_key): + """Always-apply load hook for the rank0-only row-parallel bias: zero on rank != 0. + + Unlike the shape-gated ``_load_hook`` (which only transforms when the loaded + shape differs from the sharded shape), the row-parallel bias keeps its full + shape -- only its *value* changes. Idempotent (zeroing zeros / keeping rank 0). + Takes plain ``rank``/``param_key`` (not a closure) so the pipeline cache can + serialize and rebuild it via the importable-hook path. + """ + key = prefix + param_key + if key in state_dict and rank != 0: + state_dict[key] = torch.zeros_like(state_dict[key]) + + +def _replicate_rowwise_bias(bn: WeightNode, rank: int) -> None: + """Keep a row-parallel linear's bias on rank 0 only; zero it on the others. + + The row-parallel output is summed across ranks by the trailing ``all_reduce``, + so a full bias present on every rank would be added ``world_size`` times. Zeroing + it on rank != 0 makes the all_reduce contribute the bias exactly once. The + parameter shape is unchanged (unlike the column-parallel bias, which is split), + so a dedicated always-apply load hook is used. + """ + new_bias = bn.tensor if rank == 0 else torch.zeros_like(bn.tensor) + pname = bn.node_key.rsplit(".", 1)[-1] + bn.submod._register_load_state_dict_pre_hook( + partial(_rowwise_bias_load_hook, rank=rank, param_key=pname) + ) + setattr(bn.submod, pname, torch.nn.Parameter(new_bias.detach().clone(), requires_grad=False)) + + _SHARDING_HINT_NAMES = frozenset( { "tp_mode", @@ -353,6 +384,10 @@ def apply(self, gm: GraphModule, dc: DistConfig, max_num_tokens: int = 0) -> int world_size=dc.tp_size, fused_weight_dims=fused, ) + else: + # Row-parallel: the trailing all_reduce would sum a full bias + # world_size times. Keep it on rank 0 only so it is added once. + _replicate_rowwise_bias(bn, dc.tp_rank) ad_logger.debug(f" sharded linear tp_mode={tp_mode}") return 1 @@ -595,6 +630,48 @@ def apply(self, gm: GraphModule, dc: DistConfig, max_num_tokens: int = 0) -> int return 1 if count > 0 else 0 +@ShardableNode.register(torch.ops.auto_deploy.torch_attention) +class AttentionSinksShardableNode(ShardableNode): + """``torch_attention`` with per-head ``sinks``: shard sinks along the head dim. + + Attention-sink models (e.g. GPT-OSS) add a learnable per-head sink scalar + (shape ``[num_heads]``). When the heads are TP-sharded (q/k/v colwise), the + ``sinks`` arg must follow the same head split, else the op sees full sinks + against the sharded head count. Standard attention (no ``sinks``) is a no-op. + + Gating is handled by the apply loop: attention-DP skips all non-MoE nodes + (so attention stays replicated and sinks stays full), and ``shard_layers`` + gates via the node's ``layer_type`` hint (default ``"mha"``). + """ + + def apply(self, gm: GraphModule, dc: DistConfig, max_num_tokens: int = 0) -> int: + if dc.tp_size <= 1: + return 0 + count = 0 + for wn in extract_weight_nodes(self.node).weights: + # Only the per-head ``sinks`` (1-D) follows the head split; never a 2-D weight. + if wn.tensor.dim() != 1: + continue + shard_weight_tensor( + gm=gm, + weight_tensor=wn.tensor, + param_key=wn.node_key, + dim=0, + rank=dc.tp_rank, + world_size=dc.tp_size, + ) + count += 1 + if count: + ad_logger.debug(" sharded attention sinks along head dim") + return 1 if count > 0 else 0 + + @classmethod + def _strip_node_hints(cls, node: Node) -> bool: + # Leave ``torch_attention`` untouched at strip time: its ``layer_type`` is a + # benign op default that downstream backend selection reads as-is. + return False + + @ShardableNode.register(*_auto_deploy_ops("torch_rmsnorm_gated", "triton_rmsnorm_gated")) class NormShardableNode(ShardableNode): """Gated RMSNorm op: shard weight parameter.""" @@ -979,6 +1056,12 @@ class IRShardingConfig(TransformConfig): default=None, description="When set, only shard nodes whose layer_type hint is in this list.", ) + simple_shard_filter: Optional[str] = Field( + default=None, + description="Comma-separated weight-name keywords (e.g. 'lm_head'). Matching linears are " + "gather-sharded (column split + all_gather) regardless of shard_layers -- used for the " + "lm_head vocab projection, which the hint-driven sharder would otherwise replicate.", + ) enable_attention_dp: bool = Field(default=False) dist_mapping: dict[str, int] = Field(default_factory=dict) dist_config: DistConfig = Field(default_factory=DistConfig) @@ -1047,46 +1130,54 @@ def _apply_simple_shard(gm: GraphModule, dc: DistConfig) -> int: for node in list(gm.graph.nodes): if not is_any_lin_op(node): continue - weight_nodes = extract_weight_nodes(node) - if not weight_nodes.weights: - continue - for wn in weight_nodes.weights: - shard_weight_tensor( - gm=gm, - weight_tensor=wn.tensor, - param_key=wn.node_key, - dim=SplitDimension.COLUMN, - rank=dc.tp_rank, - world_size=dc.tp_size, - ) - for bn in weight_nodes.biases: - shard_weight_tensor( - gm=gm, - weight_tensor=bn.tensor, - param_key=bn.node_key, - dim=SplitDimension.COLUMN, - rank=dc.tp_rank, - world_size=dc.tp_size, - ) - enable_sharding = ShardableNode.from_node(node) - if isinstance(enable_sharding, LinearShardableNode): - enable_sharding._shard_scales( - gm, dc, weight_nodes, dim=SplitDimension.COLUMN, min_shape=1, fused=None - ) - # torch_dist_all_gather is the demollm backend op; signature is - # (tensor, dim=0, sizes=None) — plain torch.distributed all_gather, - # no strategy or symm_mem support (use the trtllm backend for those). - with gm.graph.inserting_after(node): - gather_node = gm.graph.call_function( - torch.ops.auto_deploy.torch_dist_all_gather.default, - args=(node, -1), - ) - node.replace_all_uses_with(gather_node) - gather_node.replace_input_with(gather_node, node) - num_updates += 1 + num_updates += _simple_shard_node(gm, node, dc) return num_updates +def _simple_shard_node(gm: GraphModule, node: Node, dc: DistConfig) -> int: + """Column-split one linear's weight/bias/scale, then all_gather (the "gather" mode). + + Used as the simple-shard fallback for every linear (``simple_shard_only``) and, + via ``simple_shard_filter``, to gather-shard specific linears like ``lm_head`` + (huge vocab projection) that the hint-driven sharder would otherwise replicate. + """ + weight_nodes = extract_weight_nodes(node) + if not weight_nodes.weights: + return 0 + for wn in weight_nodes.weights: + shard_weight_tensor( + gm=gm, + weight_tensor=wn.tensor, + param_key=wn.node_key, + dim=SplitDimension.COLUMN, + rank=dc.tp_rank, + world_size=dc.tp_size, + ) + for bn in weight_nodes.biases: + shard_weight_tensor( + gm=gm, + weight_tensor=bn.tensor, + param_key=bn.node_key, + dim=SplitDimension.COLUMN, + rank=dc.tp_rank, + world_size=dc.tp_size, + ) + sn = ShardableNode.from_node(node) + if isinstance(sn, LinearShardableNode): + sn._shard_scales(gm, dc, weight_nodes, dim=SplitDimension.COLUMN, min_shape=1, fused=None) + # torch_dist_all_gather is the demollm backend op; signature is + # (tensor, dim=0, sizes=None) — plain torch.distributed all_gather, + # no strategy or symm_mem support (use the trtllm backend for those). + with gm.graph.inserting_after(node): + gather_node = gm.graph.call_function( + torch.ops.auto_deploy.torch_dist_all_gather.default, + args=(node, -1), + ) + node.replace_all_uses_with(gather_node) + gather_node.replace_input_with(gather_node, node) + return 1 + + # ============================================================================= # Transform classes # ============================================================================= @@ -1190,6 +1281,8 @@ def _apply( _log_sharding_result(dc, num_updates) else: shard_layers = self.config.shard_layers + ssf = self.config.simple_shard_filter + simple_shard_filter = [k.strip() for k in ssf.split(",") if k.strip()] if ssf else None num_skipped = 0 all_dead_nodes = [] @@ -1202,6 +1295,14 @@ def _apply( shardable_node, (MoEShardableNode, StackedMoEShardableNode) ): continue + # Gather-shard simple_shard_filter matches (e.g. lm_head) regardless of + # shard_layers: column split + all_gather of the huge vocab projection. + if simple_shard_filter and is_any_lin_op(node): + wnodes = extract_weight_nodes(node) + key = wnodes.weights[0].node_key if wnodes.weights else "" + if any(kw in key for kw in simple_shard_filter): + num_updates += _simple_shard_node(gm, node, dc) + continue if shard_layers is not None: [lt] = extract_op_args(node, "layer_type") if lt is not None and lt not in shard_layers: diff --git a/tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py b/tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py index f18e09d3453e..b0f827ead7d7 100644 --- a/tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py +++ b/tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_ir_equivalence.py @@ -86,9 +86,13 @@ def _has_ir_markers(gm) -> bool: This helper is the in-test equivalent of the follow-up PR's ``has_sharding_ir_markers`` dispatcher in ``sharding_ir.py``. """ - target = torch.ops.auto_deploy.all_reduce + # Match both the OpOverloadPacket and its ``.default`` overload: torch.export + # emits the overload (``all_reduce.default``) as node.target, which is NOT == + # the packet, so comparing against the packet alone silently misses every + # marker and skips sharding for all IR families. + targets = (torch.ops.auto_deploy.all_reduce, torch.ops.auto_deploy.all_reduce.default) for node in gm.graph.nodes: - if node.op == "call_function" and node.target == target: + if node.op == "call_function" and node.target in targets: return True return False @@ -121,29 +125,38 @@ def _has_ir_markers(gm) -> bool: WEIGHT_SEED = 0 INPUT_SEED = 42 -# bf16 is the default forward dtype for unquantized AutoDeploy deployments. -# fp32 would give tighter numerics but production runs in bf16, so that's -# what the equivalence test should validate. With random init std=0.05 and a -# deterministic-router fix applied to MoE blocks, clean sharding produces -# rel_rmse < 0.012 on every IR family; sabotaged sharding produces > 0.05. -FORWARD_DTYPE = torch.bfloat16 - -# Random init std. Small enough that 4 stacked layers don't blow up in bf16, -# large enough that the per-rank contribution missing under sabotage is -# detectable. Anything below ~0.03 makes sabotage indistinguishable from -# noise on dense models; anything above ~0.1 starts triggering bf16 routing -# noise in MoE blocks even with the deterministic-router fix. +# fp32 forward dtype. This test validates the *sharding transform* (weight / scale +# / bias slicing + collective insertion), which is dtype-independent math, so fp32 +# isolates it cleanly: correct sharding reproduces the unsharded output to ~1e-6, +# while a broken shard (missing collective, double-counted row-parallel bias, ...) +# diverges by O(1). bf16 instead conflates real bugs with precision noise (up to +# ~0.05 on attention-heavy models such as gpt_oss/step3p7, overlapping the sabotage +# band). bf16 production numerics are covered by the accuracy tests, not here. +FORWARD_DTYPE = torch.float32 + +# Random init std. Large enough that the per-rank contribution missing under +# sabotage is clearly detectable, small enough that 4 stacked layers stay +# numerically sane. INIT_STD = 0.05 -# Relative-RMSE tolerance: ``||y_s - y_u||_F / ||y_u||_F``. Scale-invariant -# across models with very different output magnitudes (dense models have -# ``|y|`` ~0.08, MoE models ~3.6). Picked to be above the worst clean -# rel_rmse observed on any IR family (~0.012 on qwen3_5_moe due to -# softmax-amplified bf16 noise in router weights) and well below the -# smallest sabotage rel_rmse (~0.05 on dense models). Override via -# ``SHARDING_IR_REL_RMSE_TOL`` env var when triaging. +# Relative-RMSE tolerance: ``||y_s - y_u||_F / ||y_u||_F``. In fp32, correct +# sharding reproduces the unsharded output to ~1e-6 while a broken shard diverges +# by O(1), so this loose bound trivially separates the two (no per-model tuning). +# Override via ``SHARDING_IR_REL_RMSE_TOL`` env var when triaging. REL_RMSE_TOL = 0.02 +# Per-family ``shard_layers`` whitelist. ``None`` (the default) shards every linear, +# which corrupts models that intend some weights replicated -- e.g. qwen3_5_moe's +# shared expert (added after the routed all_reduce, with no all_reduce of its own) -- +# so those layer_types are excluded here. ``lm_head`` is excluded for both families: +# its vocab-parallel gather (colwise + all_gather) is applied by a separate transform +# this offline harness does not run, so lm_head is replicated here and validated +# elsewhere. Keyed by modeling short name (``modeling_.py`` -> ````). +_SHARD_LAYERS_BY_FAMILY = { + "qwen3_5_moe": ["moe", "delta", "mha"], + "gpt_oss": ["mha", "moe"], +} + pytestmark = pytest.mark.threadleak(enabled=False) @@ -340,8 +353,16 @@ def _run_equivalence_job_impl( moe_ep_size=dist_cfg_spec["moe_ep_size"], enable_attention_dp=enable_attention_dp, ) + apply_hints_cfg = {"stage": "sharding", "enabled": True} + family = Path(modeling_file).stem.removeprefix("modeling_") + shard_layers = _SHARD_LAYERS_BY_FAMILY.get(family) + if shard_layers is not None: + # Mirror the model's production recipe: shard only these layer_types and + # leave the rest (e.g. qwen3_5_moe's untagged shared expert / lm_head) + # replicated, matching how the model is actually deployed. + apply_hints_cfg["shard_layers"] = shard_layers sharded_transforms = { - "apply_sharding_hints": {"stage": "sharding", "enabled": True}, + "apply_sharding_hints": apply_hints_cfg, "strip_sharding_hints": {"stage": "weight_load"}, } optimizer = InferenceOptimizer(factory=None, config=sharded_transforms, dist_config=dist_config)