sglang 0.2.13__py3-none-any.whl → 0.2.14.post1__py3-none-any.whl

sglang/srt/hf_transformers_utils.py

@@ -30,14 +30,19 @@ from transformers import (
     PreTrainedTokenizer,
     PreTrainedTokenizerFast,
 )
-from vllm.transformers_utils.configs import ChatGLMConfig, DbrxConfig
 
-from sglang.srt.utils import is_multimodal_model
+try:
+    from vllm.transformers_utils.configs import ChatGLMConfig, DbrxConfig
+
+    _CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = {
+        ChatGLMConfig.model_type: ChatGLMConfig,
+        DbrxConfig.model_type: DbrxConfig,
+    }
+except ImportError:
+    # We want this file to run without vllm dependency
+    _CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = {}
 
-_CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = {
-    ChatGLMConfig.model_type: ChatGLMConfig,
-    DbrxConfig.model_type: DbrxConfig,
-}
+from sglang.srt.utils import is_multimodal_model
 
 
 def download_from_hf(model_path: str):
@@ -137,18 +142,6 @@ def get_tokenizer(
             raise ValueError("Cannot use the fast tokenizer in slow tokenizer mode.")
         kwargs["use_fast"] = False
 
-    if (
-        "llama" in tokenizer_name.lower()
-        and kwargs.get("use_fast", True)
-        and tokenizer_name != _FAST_LLAMA_TOKENIZER
-    ):
-        pass
-        # warnings.warn(
-        #    "For some LLaMA V1 models, initializing the fast tokenizer may "
-        #    "take a long time. To reduce the initialization time, consider "
-        #    f"using '{_FAST_LLAMA_TOKENIZER}' instead of the original "
-        #    "tokenizer."
-        # )
     try:
         tokenizer = AutoTokenizer.from_pretrained(
             tokenizer_name,
@@ -229,6 +222,8 @@ class TiktokenTokenizer:
         }
         assert tok_dict["word_split"] == "V1"
 
+        default_allowed_special = None
+
         kwargs = {
             "name": name,
             "pat_str": tok_dict.get("pat_str", PAT_STR_B),
@@ -242,14 +237,18 @@ class TiktokenTokenizer:
                     for bytes_list in tok_dict["default_allowed_special"]
                 ]
             )
-        else:
-            default_allowed_special = None
         if "vocab_size" in tok_dict:
             kwargs["explicit_n_vocab"] = tok_dict["vocab_size"]
 
+        PAD = "<|pad|>"
+        EOS = "<|eos|>"
+        SEP = "<|separator|>"
+
+        DEFAULT_CONTROL_TOKENS = {"pad": PAD, "sep": EOS, "eos": SEP}
+
         tokenizer = tiktoken.Encoding(**kwargs)
         tokenizer._default_allowed_special = default_allowed_special or set()
-        tokenizer._default_allowed_special |= {"<|separator|>"}
+        tokenizer._control_tokens = DEFAULT_CONTROL_TOKENS
 
         def encode_patched(
             self,
@@ -266,14 +265,14 @@ class TiktokenTokenizer:
                 self,
                 text,
                 allowed_special=allowed_special,
-                disallowed_special=disallowed_special,
+                disallowed_special=(),
             )
 
         tokenizer.encode = functools.partial(encode_patched, tokenizer)
 
         # Convert to HF interface
         self.tokenizer = tokenizer
-        self.eos_token_id = tokenizer._special_tokens["<|eos|>"]
+        self.eos_token_id = tokenizer._special_tokens[EOS]
         self.vocab_size = tokenizer.n_vocab
         self.chat_template = Template(
             "{% for message in messages %}{% if message['role'] == 'user' %}{{ 'Human: ' + message['content'].strip() + '<|separator|>\n\n' }}{% elif message['role'] == 'system' %}{{ 'System: ' + message['content'].strip() + '<|separator|>\n\n' }}{% elif message['role'] == 'assistant' %}{{ 'Assistant: '  + message['content'] + '<|separator|>\n\n' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ 'Assistant:' }}{% endif %}"

sglang/srt/layers/activation.py

@@ -13,10 +13,20 @@ limitations under the License.
 
 """Fused operators for activation layers."""
 
+from typing import Optional
+
 import torch
+import torch.nn as nn
 import torch.nn.functional as F
-from flashinfer.activation import silu_and_mul
+from flashinfer.activation import gelu_tanh_and_mul, silu_and_mul
+from vllm.distributed import (
+    divide,
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+)
 from vllm.model_executor.custom_op import CustomOp
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.utils import set_weight_attrs
 
 
 class SiluAndMul(CustomOp):
@@ -30,3 +40,92 @@ class SiluAndMul(CustomOp):
         out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
         silu_and_mul(x, out)
         return out
+
+
+class GeluAndMul(CustomOp):
+    def __init__(self, **kwargs):
+        super().__init__()
+
+    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
+        d = x.shape[-1] // 2
+        return F.gelu(x[..., :d], approximate="tanh") * x[..., d:]
+
+    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
+        d = x.shape[-1] // 2
+        output_shape = x.shape[:-1] + (d,)
+        out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
+        gelu_tanh_and_mul(x, out)
+        return out
+
+
+class ScaledActivation(nn.Module):
+    """An activation function with post-scale parameters.
+
+    This is used for some quantization methods like AWQ.
+    """
+
+    def __init__(
+        self,
+        act_module: nn.Module,
+        intermediate_size: int,
+        input_is_parallel: bool = True,
+        params_dtype: Optional[torch.dtype] = None,
+    ):
+        super().__init__()
+        self.act = act_module
+        self.input_is_parallel = input_is_parallel
+        if input_is_parallel:
+            tp_size = get_tensor_model_parallel_world_size()
+            intermediate_size_per_partition = divide(intermediate_size, tp_size)
+        else:
+            intermediate_size_per_partition = intermediate_size
+        if params_dtype is None:
+            params_dtype = torch.get_default_dtype()
+        self.scales = nn.Parameter(
+            torch.empty(intermediate_size_per_partition, dtype=params_dtype)
+        )
+        set_weight_attrs(self.scales, {"weight_loader": self.weight_loader})
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.act(x) / self.scales
+
+    def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
+        param_data = param.data
+        if self.input_is_parallel:
+            tp_rank = get_tensor_model_parallel_rank()
+            shard_size = param_data.shape[0]
+            start_idx = tp_rank * shard_size
+            loaded_weight = loaded_weight.narrow(0, start_idx, shard_size)
+        assert param_data.shape == loaded_weight.shape
+        param_data.copy_(loaded_weight)
+
+
+_ACTIVATION_REGISTRY = {
+    "gelu": nn.GELU(),
+    "gelu_pytorch_tanh": nn.GELU(approximate="tanh"),
+}
+
+
+def get_act_fn(
+    act_fn_name: str,
+    quant_config: Optional[QuantizationConfig] = None,
+    intermediate_size: Optional[int] = None,
+    input_is_parallel: bool = True,
+    params_dtype: Optional[torch.dtype] = None,
+) -> nn.Module:
+    """Get an activation function by name."""
+    act_fn_name = act_fn_name.lower()
+    if act_fn_name not in _ACTIVATION_REGISTRY:
+        raise ValueError(f"Activation function {act_fn_name!r} is not supported.")
+
+    act_fn = _ACTIVATION_REGISTRY[act_fn_name]
+    if quant_config is not None and act_fn_name in quant_config.get_scaled_act_names():
+        if intermediate_size is None:
+            raise ValueError(
+                "intermediate_size must be specified for scaled "
+                "activation functions."
+            )
+        return ScaledActivation(
+            act_fn, intermediate_size, input_is_parallel, params_dtype
+        )
+    return act_fn

sglang/srt/layers/decode_attention.py

@@ -26,7 +26,7 @@ import triton.language as tl
 
 from sglang.srt.managers.schedule_batch import global_server_args_dict
 
-if global_server_args_dict.get("attention_reduce_in_fp32", False):
+if global_server_args_dict.get("triton_attention_reduce_in_fp32", False):
     REDUCE_TRITON_TYPE = tl.float32
     REDUCE_TORCH_TYPE = torch.float32
 else:
@@ -58,7 +58,6 @@ def _fwd_kernel_stage1(
     att_stride_h,
     kv_group_num: tl.constexpr,
     BLOCK_DMODEL: tl.constexpr,
-    BLOCK_DPE: tl.constexpr,
     BLOCK_N: tl.constexpr,
     logit_cap: tl.constexpr,
 ):
@@ -78,10 +77,6 @@ def _fwd_kernel_stage1(
 
     off_q = cur_batch * stride_qbs + cur_head * stride_qh + offs_d
 
-    if BLOCK_DPE > 0:
-        offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE)
-        off_qpe = cur_batch * stride_qbs + cur_head * stride_qh + offs_dpe
-
     offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N)
 
     block_stard_index = start_n * BLOCK_N
@@ -106,19 +101,6 @@ def _fwd_kernel_stage1(
             other=0.0,
         ).to(REDUCE_TRITON_TYPE)
         att_value = tl.sum(q[None, :] * k, 1)
-        if BLOCK_DPE > 0:
-            qpe = tl.load(Q + off_qpe + start_mark).to(REDUCE_TRITON_TYPE)
-            offs_buf_kpe = (
-                k_loc[:, None] * stride_buf_kbs
-                + cur_kv_head * stride_buf_kh
-                + offs_dpe[None, :]
-            )
-            kpe = tl.load(
-                K_Buffer + offs_buf_kpe,
-                mask=offs_n_new[:, None] < cur_batch_end_index,
-                other=0.0,
-            ).to(REDUCE_TRITON_TYPE)
-            att_value += tl.sum(qpe[None, :] * kpe, 1)
         att_value *= sm_scale
 
         if logit_cap > 0:
@@ -214,14 +196,7 @@ def _decode_att_m_fwd(
     # shape constraints
     Lq, Lk = q.shape[-1], k_buffer.shape[-1]
     assert Lq == Lk
-    assert Lk in {16, 32, 64, 128, 256, 576}
-
-    if Lk == 576:
-        BLOCK_DMODEL = 512
-        BLOCK_DPE = 64
-    else:
-        BLOCK_DMODEL = Lk
-        BLOCK_DPE = 0
+    assert Lk in {16, 32, 64, 128, 256}
 
     batch, head_num = B_req_idx.shape[0], q.shape[1]
 
@@ -249,8 +224,7 @@ def _decode_att_m_fwd(
         k_buffer.stride(1),
         att_out.stride(0),
         kv_group_num=kv_group_num,
-        BLOCK_DMODEL=BLOCK_DMODEL,
-        BLOCK_DPE=BLOCK_DPE,
+        BLOCK_DMODEL=Lk,
         BLOCK_N=BLOCK,
         logit_cap=logit_cap,
         num_warps=num_warps,
@@ -296,6 +270,293 @@ def _decode_softmax_reducev_fwd(
     )
 
 
+@triton.jit
+def _fwd_grouped_kernel_stage1(
+    Q,
+    K_Buffer,
+    sm_scale,
+    Req_to_tokens,
+    B_req_idx,
+    B_Start_Loc,
+    B_Seqlen,
+    Att_Out,
+    stride_req_to_tokens_b,
+    stride_qbs,
+    stride_qh,
+    stride_buf_kbs,
+    stride_buf_kh,
+    att_stride_h,
+    kv_group_num: tl.constexpr,
+    q_head_num: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_DPE: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    logit_cap: tl.constexpr,
+):
+    cur_batch = tl.program_id(0)
+    cur_kv_head = tl.program_id(1)
+    start_n = tl.program_id(2)
+
+    cur_head = cur_kv_head * kv_group_num + tl.arange(0, BLOCK_H)
+    mask_h = cur_head < (cur_kv_head + 1) * kv_group_num
+    mask_h = mask_h & (cur_head < q_head_num)
+
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
+    cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
+    cur_batch_req_idx = tl.load(B_req_idx + cur_batch)
+
+    cur_batch_start_index = 0
+    cur_batch_end_index = cur_batch_seq_len
+
+    offs_q = cur_batch * stride_qbs + cur_head[:, None] * stride_qh + offs_d[None, :]
+
+    if BLOCK_DPE > 0:
+        offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE)
+        off_qpe = (
+            cur_batch * stride_qbs + cur_head[:, None] * stride_qh + offs_dpe[None, :]
+        )
+
+    offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    block_stard_index = start_n * BLOCK_N
+    block_mask = tl.where(block_stard_index < cur_batch_seq_len, 1, 0)
+
+    for start_mark in range(0, block_mask, 1):
+        q = tl.load(Q + offs_q + start_mark, mask=mask_h[:, None]).to(
+            REDUCE_TRITON_TYPE
+        )
+        offs_n_new = cur_batch_start_index + offs_n
+        k_loc = tl.load(
+            Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + offs_n_new,
+            mask=offs_n_new < cur_batch_end_index,
+            other=0,
+        )
+        offs_buf_k = (
+            k_loc[None, :] * stride_buf_kbs
+            + cur_kv_head * stride_buf_kh
+            + offs_d[:, None]
+        )
+        k = tl.load(
+            K_Buffer + offs_buf_k,
+            mask=offs_n_new[None, :] < cur_batch_end_index,
+            other=0.0,
+        ).to(REDUCE_TRITON_TYPE)
+        qk = tl.dot(q, k)
+        if BLOCK_DPE > 0:
+            qpe = tl.load(Q + off_qpe + start_mark, mask=mask_h[:, None]).to(
+                REDUCE_TRITON_TYPE
+            )
+            offs_buf_kpe = (
+                k_loc[None, :] * stride_buf_kbs
+                + cur_kv_head * stride_buf_kh
+                + offs_dpe[:, None]
+            )
+            kpe = tl.load(
+                K_Buffer + offs_buf_kpe,
+                mask=offs_n_new[None, :] < cur_batch_end_index,
+                other=0.0,
+            ).to(REDUCE_TRITON_TYPE)
+            qk += tl.dot(qpe, kpe)
+        qk *= sm_scale
+
+        if logit_cap > 0:
+            qk = logit_cap * tanh(qk / logit_cap)
+
+        offs_o = cur_head[:, None] * att_stride_h + (
+            cur_batch_in_all_start_index + offs_n[None, :]
+        )
+
+        tl.store(
+            Att_Out + offs_o,
+            qk,
+            mask=mask_h[:, None] & (offs_n_new[None, :] < cur_batch_end_index),
+        )
+
+
+@triton.jit
+def _fwd_grouped_kernel_stage2(
+    Logics,
+    V_Buffer,
+    Out,
+    Req_to_tokens,
+    B_req_idx,
+    B_Start_Loc,
+    B_Seqlen,
+    stride_logic_h,
+    stride_buf_vbs,
+    stride_buf_vh,
+    stride_obs,
+    stride_oh,
+    stride_req_to_token_b,
+    kv_group_num: tl.constexpr,
+    q_head_num: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+):
+    cur_batch = tl.program_id(0)
+    cur_kv_head = tl.program_id(1)
+
+    cur_head = cur_kv_head * kv_group_num + tl.arange(0, BLOCK_H)
+    mask_h = cur_head < (cur_kv_head + 1) * kv_group_num
+    mask_h = mask_h & (cur_head < q_head_num)
+
+    cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
+    cur_batch_start_loc = tl.load(B_Start_Loc + cur_batch)
+    cur_batch_req_idx = tl.load(B_req_idx + cur_batch)
+
+    offs_n = tl.arange(0, BLOCK_N)
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+
+    offs_buf_v = cur_kv_head * stride_buf_vh + offs_d[None, :]
+    v_ptrs = V_Buffer + offs_buf_v
+
+    e_max = tl.zeros([BLOCK_H], dtype=tl.float32) - float("inf")
+    e_sum = tl.zeros([BLOCK_H], dtype=tl.float32)
+    acc = tl.zeros([BLOCK_H, BLOCK_DMODEL], dtype=tl.float32)
+
+    for start_n in range(0, cur_batch_seq_len, BLOCK_N):
+        start_n = tl.multiple_of(start_n, BLOCK_N)
+        v_index = tl.load(
+            Req_to_tokens
+            + cur_batch_req_idx * stride_req_to_token_b
+            + (start_n + offs_n),
+            mask=(start_n + offs_n) < cur_batch_seq_len,
+            other=0,
+        )
+
+        offs_qk = cur_head[:, None] * stride_logic_h + (
+            cur_batch_start_loc + start_n + offs_n[None, :]
+        )
+
+        qk = tl.load(
+            Logics + offs_qk,
+            mask=mask_h[:, None] & (start_n + offs_n[None, :] < cur_batch_seq_len),
+            other=float("-inf"),
+        )
+
+        n_e_max = tl.maximum(tl.max(qk, 1), e_max)
+        old_scale = tl.exp(e_max - n_e_max)
+        p = tl.exp(qk - n_e_max[:, None])
+        e_sum = e_sum * old_scale + tl.sum(p, 1)
+        v = tl.load(v_ptrs + v_index[:, None] * stride_buf_vbs)
+        p = p.to(v.dtype)
+        acc = acc * old_scale[:, None] + tl.dot(p, v)
+        e_max = n_e_max
+
+    acc = acc / e_sum[:, None]
+    off_o = cur_batch * stride_obs + cur_head[:, None] * stride_oh + offs_d[None, :]
+    out_ptrs = Out + off_o
+    tl.store(out_ptrs, acc, mask=mask_h[:, None])
+
+
+def _decode_grouped_att_m_fwd(
+    q,
+    k_buffer,
+    att_out,
+    Req_to_tokens,
+    B_req_idx,
+    B_Start_Loc,
+    B_Seqlen,
+    max_len_in_batch,
+    sm_scale,
+    logit_cap,
+):
+    BLOCK = 32
+    # shape constraints
+    Lq, Lk = q.shape[-1], k_buffer.shape[-1]
+    assert Lq == Lk
+    assert Lk in {16, 32, 64, 128, 256, 576}
+
+    if Lk == 576:
+        BLOCK_DMODEL = 512
+        BLOCK_DPE = 64
+    else:
+        BLOCK_DMODEL = Lk
+        BLOCK_DPE = 0
+
+    batch, head_num = B_req_idx.shape[0], q.shape[1]
+    kv_group_num = q.shape[1] // k_buffer.shape[1]
+
+    BLOCK_H = max(16, triton.next_power_of_2(kv_group_num))
+    grid = (
+        batch,
+        triton.cdiv(head_num, min(BLOCK_H, kv_group_num)),
+        triton.cdiv(max_len_in_batch, BLOCK),
+    )
+
+    num_warps = 4
+
+    _fwd_grouped_kernel_stage1[grid](
+        q,
+        k_buffer,
+        sm_scale,
+        Req_to_tokens,
+        B_req_idx,
+        B_Start_Loc,
+        B_Seqlen,
+        att_out,
+        Req_to_tokens.stride(0),
+        q.stride(0),
+        q.stride(1),
+        k_buffer.stride(0),
+        k_buffer.stride(1),
+        att_out.stride(0),
+        kv_group_num=kv_group_num,
+        q_head_num=head_num,
+        BLOCK_DMODEL=BLOCK_DMODEL,
+        BLOCK_DPE=BLOCK_DPE,
+        BLOCK_N=BLOCK,
+        BLOCK_H=BLOCK_H,
+        logit_cap=logit_cap,
+        num_warps=num_warps,
+        num_stages=1,
+    )
+
+
+def _decode_grouped_softmax_reducev_fwd(
+    logics,
+    v_buffer,
+    o,
+    req_to_tokens,
+    b_req_idx,
+    b_start_loc,
+    b_seq_len,
+):
+    BLOCK = 128
+    batch, head_num = b_seq_len.shape[0], logics.shape[0]
+    kv_group_num = logics.shape[0] // v_buffer.shape[1]
+    BLOCK_H = max(16, triton.next_power_of_2(kv_group_num))
+    grid = (batch, triton.cdiv(head_num, min(BLOCK_H, kv_group_num)), 1)
+
+    num_warps = 8
+
+    _fwd_grouped_kernel_stage2[grid](
+        logics,
+        v_buffer,
+        o,
+        req_to_tokens,
+        b_req_idx,
+        b_start_loc,
+        b_seq_len,
+        logics.stride(0),
+        v_buffer.stride(0),
+        v_buffer.stride(1),
+        o.stride(0),
+        o.stride(1),
+        req_to_tokens.stride(0),
+        kv_group_num=kv_group_num,
+        q_head_num=head_num,
+        BLOCK_DMODEL=v_buffer.shape[-1],
+        BLOCK_N=BLOCK,
+        BLOCK_H=BLOCK_H,
+        num_warps=num_warps,
+        num_stages=1,
+    )
+
+
 def decode_attention_fwd(
     q,
     k_buffer,
@@ -316,24 +577,51 @@ def decode_attention_fwd(
             (q.shape[-2], total_num_tokens), dtype=REDUCE_TORCH_TYPE, device="cuda"
         )
 
-    _decode_att_m_fwd(
-        q,
-        k_buffer,
-        att_m,
-        req_to_token,
-        b_req_idx,
-        b_start_loc,
-        b_seq_len,
-        max_len_in_batch,
-        sm_scale,
-        logit_cap,
-    )
-    _decode_softmax_reducev_fwd(
-        att_m,
-        v_buffer,
-        o,
-        req_to_token,
-        b_req_idx,
-        b_start_loc,
-        b_seq_len,
-    )
+    kv_group_num = q.shape[1] // v_buffer.shape[1]
+
+    if kv_group_num == 1:
+        # MHA
+        _decode_att_m_fwd(
+            q,
+            k_buffer,
+            att_m,
+            req_to_token,
+            b_req_idx,
+            b_start_loc,
+            b_seq_len,
+            max_len_in_batch,
+            sm_scale,
+            logit_cap,
+        )
+        _decode_softmax_reducev_fwd(
+            att_m,
+            v_buffer,
+            o,
+            req_to_token,
+            b_req_idx,
+            b_start_loc,
+            b_seq_len,
+        )
+    else:
+        # GQA/MQA/MLA
+        _decode_grouped_att_m_fwd(
+            q,
+            k_buffer,
+            att_m,
+            req_to_token,
+            b_req_idx,
+            b_start_loc,
+            b_seq_len,
+            max_len_in_batch,
+            sm_scale,
+            logit_cap,
+        )
+        _decode_grouped_softmax_reducev_fwd(
+            att_m,
+            v_buffer,
+            o,
+            req_to_token,
+            b_req_idx,
+            b_start_loc,
+            b_seq_len,
+        )

sglang/srt/layers/fused_moe/layer.py

@@ -239,7 +239,7 @@ class FusedMoE(torch.nn.Module):
         weight_name: str,
         shard_id: int,
         expert_id: int,
-        pre_sharded: bool,
+        use_presharded_weights: bool = False,
     ):
         param_data = param.data
 
@@ -273,7 +273,7 @@ class FusedMoE(torch.nn.Module):
         else:
             tp_rank = get_tensor_model_parallel_rank()
             shard_size = self.intermediate_size_per_partition
-            if pre_sharded:
+            if use_presharded_weights:
                 shard = slice(None)
             else:
                 shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size)