Function compute_weight_and_optimizer_memory

linear_moe/utils.py:22–132

(args, verbose=False)

    return wrapper

def compute_weight_and_optimizer_memory(args, verbose=False):
    # Attention projection size.
    query_projection_size = args.kv_channels * args.num_attention_heads
    query_projection_to_hidden_size_ratio = query_projection_size / args.hidden_size
    # Group Query Attention.
    if not args.group_query_attention:
        args.num_query_groups = args.num_attention_heads
    # MoE.
    num_experts = 1 if args.num_experts is None else args.num_experts
    gated_linear_multiplier = 3 / 2 if args.swiglu else 1

    # Fine-grained Experts + Shared Experts
    if args.moe_ffn_hidden_size is not None:
        intermediate_hidden_size = args.moe_ffn_hidden_size
    else:
        intermediate_hidden_size = args.ffn_hidden_size
    
    shared_moe_ffn_hidden_size = args.shared_moe_ffn_hidden_size if args.shared_moe_ffn_hidden_size else 0

    num_parameters_in_transformer_layers = (
        2
        * args.num_layers
        * args.hidden_size
        * args.hidden_size
        * (
            # Attention.
            (
                (1 + (args.num_query_groups / args.num_attention_heads))
                * query_projection_to_hidden_size_ratio
            )
            # MLP.
            + ((intermediate_hidden_size / args.hidden_size) * num_experts * gated_linear_multiplier + (shared_moe_ffn_hidden_size / args.hidden_size) * gated_linear_multiplier)
            # Transformer layernorms.
            + (2 / args.hidden_size)
            # Final layernorm.
            + (1 / (args.num_layers * args.hidden_size))
        )
    )
    num_parameters_in_transformer_moe_act_layers = (
        2
        * args.num_layers
        * args.hidden_size
        * args.hidden_size
        * (
            # Attention.
            (
                (1 + (args.num_query_groups / args.num_attention_heads))
                * query_projection_to_hidden_size_ratio
            )
            # MLP.
            + ((intermediate_hidden_size / args.hidden_size) * args.moe_router_topk * gated_linear_multiplier + (shared_moe_ffn_hidden_size / args.hidden_size) * gated_linear_multiplier)
            # Transformer layernorms.
            + (2 / args.hidden_size)
            # Final layernorm.
            + (1 / (args.num_layers * args.hidden_size))
        )
    )