Repository: oumi-ai/oumi
File: scripts/memcalc.py (Python)

Code before:
    """Gets the total number of bytes used by the data batch."""
    batch_size = config.training.per_device_train_batch_size
    model_max_length = get_seq_len(config, model_config)
    return batch_size * model_max_length * get_bytes_per_unit(_TOKEN_DTYPE)


# TODO: Find a static way to calculate this
def get_model_bytes(model: torch.nn.Module, bytes_per_unit: int) -> int:
    """Gets the total number of bytes used by the loaded model."""
    num_total_params = count_model_parameters(model)
    return num_total_params.all_params * bytes_per_unit


# TODO: Support more optimizers
def get_optim_bytes(config: TrainingConfig, model_bytes: int) -> int:
    """Gets the total number of bytes used by the optimizer."""
    optim = config.training.optimizer
    if optim in ["adamw_torch", "adamw_torch_fused"]:
        multiplier = 2
    elif optim == "adafactor":
        multiplier = 0.3
    elif optim == "sgd":
        # TODO: Account for momentum
        multiplier = 0
    else:
        raise ValueError(f"Unsupported optimizer: {optim}")
    print(f"Optimizer {optim} uses {multiplier}x model bytes")
    return int(model_bytes * multiplier)


def get_gradient_bytes(config: TrainingConfig, model_bytes: int) -> int:
    """Gets the total number of bytes used by gradients."""
    if config.training.gradient_accumulation_steps > 1:
        return model_bytes * 2
    return model_bytes


def get_activation_bytes(
    config: TrainingConfig, model_config: ModelConfig, data_bytes: int
) -> int:
    """Gets the total number of bytes used by activations."""
    vocab_size = model_config.vocab_size
    num_layers = model_config.num_layers
    hidden_dim = model_config.hidden_dim
    seq_len = get_seq_len(config, model_config)
    num_kv_heads = model_config.num_kv_heads
    # TODO: Verify this formula
    embedding_bytes = vocab_size
    lm_head_bytes = vocab_size
    transformer_bytes = num_layers * (14 * hidden_dim + seq_len * num_kv_heads)
    total = embedding_bytes + lm_head_bytes + transformer_bytes

Review comment (performance):
Should we have a some of the core logic as a library under oumi.performance? The script would just call that.

In the future we could also have an `oumi train -c config --dry-run` that uses it

Code after:
    if optim in ["adamw_torch", "adamw_torch_fused"]:
        multiplier = 2
    elif optim == "adafactor":
        multiplier = 0.3
    elif optim == "sgd":
        # TODO: Account for momentum
        multiplier = 0
    else:
        raise ValueError(f"Unsupported optimizer: {optim}")
    print(f"Optimizer {optim} uses {multiplier}x model bytes")
    return int(model_bytes * multiplier)


def get_gradient_bytes(config: TrainingConfig, model_bytes: int) -> int:
    """Gets the total number of bytes used by gradients."""
    print("- The size of the gradient is the same as that for model weights.")
    if config.training.gradient_accumulation_steps > 1:
        print(
            "- If gradient accumulation is used, a buffer to store gradients is "
            "required. The buffer is the same size as the gradients."
        )
        return model_bytes * 2
    return model_bytes


def get_activation_bytes(
    config: TrainingConfig, model_config: ModelConfig, bytes_per_unit: int
) -> int:
    """Gets the total number of bytes used by activations."""
    vocab_size = model_config.vocab_size
    num_layers = model_config.num_layers
    hidden_dim = model_config.hidden_dim
    seq_len = get_seq_len(config, model_config)
    num_kv_heads = model_config.num_kv_heads
    # TODO: Verify this formula
    embedding_bytes = vocab_size
    lm_head_bytes = vocab_size
    transformer_bytes = num_layers * (14 * hidden_dim + seq_len * num_kv_heads)
    total = embedding_bytes + lm_head_bytes + transformer_bytes

    batch_size = config.training.per_device_train_batch_size
    model_max_length = get_seq_len(config, model_config)
    return total * batch_size * model_max_length * bytes_per_unit


def main() -> None:
    """Main function."""
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-c",
        "--config",