Function train

llm_analysis/analysis.py:2354–2503

Entry point function of training analysis for the command line interface. This uses pre-defined name-to-configuration mapping and common arguments to construct LLMAnalysis. Args: model_name (str, optional): model name to query the pre-defined `model_configs` dict, or model confi

(
    model_name="facebook_opt-1.3b",
    gpu_name="a100-sxm-40gb",
    dtype_name="w16a16e16",
    log_level="INFO",
    batch_size_per_gpu: int = None,
    gradient_accumulation_steps: int = None,
    global_batch_size: int = None,
    seq_len: int = None,
    total_num_tokens: int = None,
    activation_recomputation: int = 0,
    ds_zero: int = 0,
    dp_size: int = None,
    tp_size: int = 1,
    pp_size: int = 1,
    sp_size: int = None,
    ep_size: int = 1,
    total_num_gpus: int = None,
    layernorm_dtype_bytes: int = BYTES_FP32,
    master_weights_dtype_bytes: int = BYTES_FP32,
    other_op_bytes: int = None,
    flash_attn: bool = True,
    softmax_dropout: bool = False,
    mlp_activation_quant_bits: int = None,
    mlp_1linear_quant_bits: int = None,
    mlp_gelu_input_quant_bits: int = None,
    mlp_2linear_quant_bits: int = None,
    mlp_recompute_gelu: bool = False,
    mlp_gated_linear_units: bool = False,
    achieved_tflops: float = None,
    flops_efficiency: float = None,
    hbm_memory_efficiency: float = HBM_MEMORY_EFFICIENCY,
    intra_node_memory_efficiency=INTRA_NODE_MEMORY_EFFICIENCY,
    inter_node_memory_efficiency=INTER_NODE_MEMORY_EFFICIENCY,
    num_gpus_per_node: int = NUM_GPUS_PER_NODE,
    output_dir: str = None,
    output_file_suffix: str = "",
)

def train(
    model_name="facebook_opt-1.3b",
    gpu_name="a100-sxm-40gb",
    dtype_name="w16a16e16",
    log_level="INFO",
    batch_size_per_gpu: int = None,
    gradient_accumulation_steps: int = None,
    global_batch_size: int = None,
    seq_len: int = None,
    total_num_tokens: int = None,
    activation_recomputation: int = 0,
    ds_zero: int = 0,
    dp_size: int = None,
    tp_size: int = 1,
    pp_size: int = 1,
    sp_size: int = None,
    ep_size: int = 1,
    total_num_gpus: int = None,
    layernorm_dtype_bytes: int = BYTES_FP32,
    master_weights_dtype_bytes: int = BYTES_FP32,
    other_op_bytes: int = None,
    flash_attn: bool = True,
    softmax_dropout: bool = False,
    mlp_activation_quant_bits: int = None,
    mlp_1linear_quant_bits: int = None,
    mlp_gelu_input_quant_bits: int = None,
    mlp_2linear_quant_bits: int = None,
    mlp_recompute_gelu: bool = False,
    mlp_gated_linear_units: bool = False,
    achieved_tflops: float = None,
    flops_efficiency: float = None,
    hbm_memory_efficiency: float = HBM_MEMORY_EFFICIENCY,
    intra_node_memory_efficiency=INTRA_NODE_MEMORY_EFFICIENCY,
    inter_node_memory_efficiency=INTER_NODE_MEMORY_EFFICIENCY,
    num_gpus_per_node: int = NUM_GPUS_PER_NODE,
    output_dir: str = None,
    output_file_suffix: str = "",
) -> dict:
    """Entry point function of training analysis for the command line interface. This
    uses pre-defined name-to-configuration mapping and common arguments to construct
    LLMAnalysis.

    Args:
        model_name (str, optional): model name to query the pre-defined `model_configs` dict, or model config json file path, if not found, query Hugging Face to construct ModelConfig. Defaults to "facebook_opt-1.3b".
        gpu_name (str, optional): gpu name to query the pre-defined `gpu_configs` dict. Defaults to "a100-sxm-40gb".
        dtype_name (str, optional): data type name to pre-defined `dtype_configs` dict. Defaults to "w16a16e16".
        log_level (str, optional): logging level. Defaults to "INFO".
        batch_size_per_gpu (int, optional): batch size per GPU (micro batch size). Defaults to None.
        gradient_accumulation_steps (int, optional): gradient accumulation steps. Defaults to None.
        global_batch_size (int, optional): global batch size. Defaults to None.
        seq_len (int, optional): sequence length. Defaults to None.
        total_num_tokens (int, optional): total number of tokens used for training. Defaults to None.
        activation_recomputation (int, optional): activation recomputation strategy. See `ActivationRecomputation`. Defaults to 0.
        ds_zero (int, optional): which DeepSpeed ZeRO stage to use. See `DSZeRO`. Defaults to 0.
        dp_size (int, optional): data parallelism size. Defaults to None.
        tp_size (int, optional): tensor parallelism size. Defaults to 1.
        pp_size (int, optional): pipeline parallelism size. Defaults to 1.
        sp_size (int, optional): sequence parallelism size. Defaults to tp_size.