def parameter_name_to_sharding_dim()

in summarize_from_feedback/query_response_model.py [0:0]

• 20 lines of code
• 9 McCabe index (conditional complexity)

def parameter_name_to_sharding_dim(name: str) -> Optional[int]:
    """
    :returns: None if all parameters are same on all shards, otherwise the dimension to
        split upon.
    """
    if name in ["embedding.weight", "position_embedding.weight", "unembedding_weights"]:
        return -1
    if name.startswith("torso.resblocks"):
        match = re.search(r"torso\.resblocks\.\d+\.(.*)", name)
        torso_part = match.group(1)
        if torso_part.startswith("ln_1.") or torso_part.startswith("ln_2."):
            return None
        if _matches_any_prefix(
            torso_part, ["attn.q_proj", "attn.k_proj", "attn.v_proj", "mlp.c_fc"]
        ):
            return -1
        if _matches_any_prefix(torso_part, ["attn.c_proj.weight", "mlp.c_proj.weight"]):
            return -2
        if _matches_any_prefix(torso_part, ["attn.c_proj.bias", "mlp.c_proj.bias"]):
            return None
        raise RuntimeError(f"Unexpected parameter name: {name}")
    if name in ["ln_f.weight", "ln_f.bias"]:
        return None
    raise RuntimeError(f"Unexpected parameter name: {name}")