Function main

benchmark.py:261–488

(
    ctx: typer.Context,  # Typer Context to grab config for --verbose and passing to WandB
    hidden_sizes: Annotated[List[int], Option(help="List of hidden sizes", show_default=False)],
    num_hidden_layers: Annotated[List[int], Option(help="List of number of hidden layers", show_default=False)],
    intermediate_sizes: Annotated[List[int], Option(help="List of intermediate sizes", show_default=False)],
    parallel_attn: Annotated[List[bool], Option(is_flag=False, help="List of parallel attention flags", show_default=False)],
    sliding_window: Annotated[List[int], Option(help="Sliding window size. -1 to disable.")] = [-1],
    global_attn_every_n_layers: Annotated[List[int], Option(help="Use global attention every `n` layers and sliding window for the rest. -1 to disable.")] = [-1],
    normalization: Annotated[List[str], Option(help="Normalization type: layernorm or triton_layernorm")] = ["layernorm"],
    compile_model: Annotated[List[bool], Option(help="Compile model")] = [True],
    masked_prediction: Annotated[List[bool], Option(help="Only pass the masked tokens through the final MLM layers")] = [True],
    model_type: Annotated[List[ModelType], Option(help="Model type: MLM or Multiple Choice")] = [ModelType.mlm],
    vocab_size: Annotated[List[int], Option(help="Vocabulary size")] = [32768],
    num_samples: Annotated[int, Option(help="Number of samples")] = 1000,
    seq_length: Annotated[int, Option(help="Sequence length")] = 512,
    batch_size: Annotated[Optional[int], Option(help="Batch size (if not provided, will be set based on model size)")] = None,
    output_file: Annotated[str, Option(help="Output file name for results")] = "benchmark_results.md",
    sleep_time: Annotated[int, Option(help="Time to sleep between each model run")] = 25,
    print_model: Annotated[bool, Option(help="Print model")] = False,
    num_workers: Annotated[int, Option(help="Number of workers")] = 8,
    skip_inference: Annotated[bool, Option(help="Skip inference")] = False,
    gpu_idx: Annotated[int, Option(help="GPU index for power measurements")] = 0,
    config: Annotated[
        Optional[Path],
        Option(
            callback=conf_callback,
            is_eager=True,
            help="Relative path to YAML config file for setting options. Passing CLI options will supersede config options.",
            case_sensitive=False,
        ),
    ] = None,
)

# fmt: off
@app.command()
def main(
    ctx: typer.Context,  # Typer Context to grab config for --verbose and passing to WandB
    hidden_sizes: Annotated[List[int], Option(help="List of hidden sizes", show_default=False)],
    num_hidden_layers: Annotated[List[int], Option(help="List of number of hidden layers", show_default=False)],
    intermediate_sizes: Annotated[List[int], Option(help="List of intermediate sizes", show_default=False)],
    parallel_attn: Annotated[List[bool], Option(is_flag=False, help="List of parallel attention flags", show_default=False)],
    sliding_window: Annotated[List[int], Option(help="Sliding window size. -1 to disable.")] = [-1],
    global_attn_every_n_layers: Annotated[List[int], Option(help="Use global attention every `n` layers and sliding window for the rest. -1 to disable.")] = [-1],
    normalization: Annotated[List[str], Option(help="Normalization type: layernorm or triton_layernorm")] = ["layernorm"],
    compile_model: Annotated[List[bool], Option(help="Compile model")] = [True],
    masked_prediction: Annotated[List[bool], Option(help="Only pass the masked tokens through the final MLM layers")] = [True],
    model_type: Annotated[List[ModelType], Option(help="Model type: MLM or Multiple Choice")] = [ModelType.mlm],
    vocab_size: Annotated[List[int], Option(help="Vocabulary size")] = [32768],
    num_samples: Annotated[int, Option(help="Number of samples")] = 1000,
    seq_length: Annotated[int, Option(help="Sequence length")] = 512,
    batch_size: Annotated[Optional[int], Option(help="Batch size (if not provided, will be set based on model size)")] = None,
    output_file: Annotated[str, Option(help="Output file name for results")] = "benchmark_results.md",
    sleep_time: Annotated[int, Option(help="Time to sleep between each model run")] = 25,
    print_model: Annotated[bool, Option(help="Print model")] = False,
    num_workers: Annotated[int, Option(help="Number of workers")] = 8,
    skip_inference: Annotated[bool, Option(help="Skip inference")] = False,
    gpu_idx: Annotated[int, Option(help="GPU index for power measurements")] = 0,
    config: Annotated[
        Optional[Path],
        Option(
            callback=conf_callback,
            is_eager=True,
            help="Relative path to YAML config file for setting options. Passing CLI options will supersede config options.",
            case_sensitive=False,
        ),
    ] = None,
):
# fmt: on
    pynvml.nvmlInit()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Determine the maximum length of the lists
    max_length = max(
        len(hidden_sizes),
        len(num_hidden_layers),
        len(intermediate_sizes),
        len(parallel_attn),
        len(vocab_size),
        len(model_type),
        len(sliding_window),
        len(global_attn_every_n_layers),
        len(normalization),
        len(compile_model),
        len(masked_prediction),
    )

    # Tile lists to match the maximum length
    hidden_sizes = tile_list_to_length(hidden_sizes, max_length)
    num_hidden_layers = tile_list_to_length(num_hidden_layers, max_length)
    intermediate_sizes = tile_list_to_length(intermediate_sizes, max_length)
    parallel_attn = tile_list_to_length(parallel_attn, max_length)
    vocab_size = tile_list_to_length(vocab_size, max_length)
    model_type = tile_list_to_length(model_type, max_length)