config: a parallel size has one spelling; a patched scope declares its own (#36621)
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
ca1d7ed8e6
commit
fd40a331bf
@@ -109,9 +109,12 @@ def mixer2_gated_norm_tensor_parallel(
|
||||
|
||||
import sglang.srt.layers.attention.mamba.mixer2_rms_norm_gated as m2
|
||||
|
||||
# Force attn-TP rank through the context (the weight loader reads it via
|
||||
# get_parallel().attn_tp_rank); avoids calling initialize_dp_attention.
|
||||
with get_parallel().override(attn_tp_rank=local_rank):
|
||||
# Force the TP topology through the context (the weight loader reads
|
||||
# get_parallel().attn_tp_rank, Mixer2RMSNormGated reads tp_size / tp_rank);
|
||||
# avoids calling initialize_dp_attention.
|
||||
with get_parallel().override(
|
||||
attn_tp_rank=local_rank, tp_size=world_size, tp_rank=local_rank
|
||||
):
|
||||
# create gated-norm with TP
|
||||
mixer = m2.Mixer2RMSNormGated(
|
||||
full_hidden_size=hidden_size,
|
||||
|
||||
Reference in New Issue
Block a user