def get_shard_fix_factor()

in summarize_from_feedback/query_response_model.py [0:0]

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

def get_shard_fix_factor(name: str, model_H: Hyperparams, old_model_H: Hyperparams) -> float:
    # Hack to fix some bugs with our sharding code
    if name.startswith("torso.resblocks"):
        match = re.search(r"torso\.resblocks\.\d+\.(.*)", name)
        torso_part = match.group(1)

        # bias is added before all-reduce, which means with more shards, the
        # weights are closer to 0 than expected
        if _matches_any_prefix(torso_part, ["attn.c_proj.bias", "mlp.c_proj.bias"]):
            return float(old_model_H.n_shards) / model_H.n_shards

        if (
            _matches_any_prefix(
                torso_part,
                [
                    "attn.q_proj.weight",
                    "attn.k_proj.weight",
                    "attn.q_proj.bias",
                    "attn.k_proj.bias",
                ],
            )
            and old_model_H.use_blocksparse_attn
        ):
            return np.sqrt(np.sqrt(float(old_model_H.n_shards) / old_model_H.heads))
    return 1.0