[tab] utils/train_utils.py (python)
import jax
import optax
import operator

from train_tokenizer import Args as TokenizerArgs
from train_lam import Args as LAMArgs
from train_dynamics import Args as DynamicsArgs

def get_lr_schedule(
    lr_schedule: str,
    init_lr: float,
    max_lr: float,
    decay_end: float,
    total_steps: int,
    warmup_steps: int,
    wsd_decay_steps: int,
) -> optax.Schedule:
    supported_schedules = ["wsd", "cos"]
    if lr_schedule == "cos":
        assert (
            warmup_steps <= total_steps
        ), "Warmup steps can't be greater than total steps."
        return optax.warmup_cosine_decay_schedule(
            init_value=init_lr,
            peak_value=max_lr,
            warmup_steps=warmup_steps,
            decay_steps=total_steps,  # Note: decay_steps includes the warmup steps, so we need to pass total value
            end_value=decay_end,
        )
    elif lr_schedule == "wsd":
        assert (
            warmup_steps + wsd_decay_steps <= total_steps
        ), "Warmup and decay period is longer than total steps."
        schedules = [
            optax.linear_schedule(
                init_value=init_lr, end_value=max_lr, transition_steps=warmup_steps
            ),
            optax.constant_schedule(value=max_lr),
            optax.linear_schedule(
                init_value=max_lr, end_value=decay_end, transition_steps=wsd_decay_steps
            ),
        ]
        boundaries = [warmup_steps, total_steps - wsd_decay_steps]
        return optax.join_schedules(schedules, boundaries)
    else:
        raise ValueError(
            f"Learning rate schedule not supported. Please use one of {supported_schedules}"
        )


def _count_component(component_params):
    """Count total parameters in a component."""
    params_sizes = jax.tree.map(jax.numpy.size, component_params)
    total_parameters = jax.tree.reduce(operator.add, params_sizes)
    return total_parameters


def count_parameters_by_component(params):
    """Count parameters for each component of the model.

    Args:
        params: Model parameters from nnx.split(model, nnx.Param, ...)

    Returns:
        Dictionary with parameter counts for each component
    """
    component_names = list(params.keys())
    print(f"Counting all components: {component_names}")

    counts = {}
    total_params = 0

    for name in component_names:
        component_params = params[name]
        count = _count_component(component_params)
        counts[name] = count
        total_params += count

    counts["total"] = total_params
    return counts


def bytes_to_gb(num_bytes):
    return num_bytes / (1024**3)


def print_compiled_memory_stats(compiled_stats):
    """from: https://github.com/AI-Hypercomputer/maxtext/blob/b18829fbaa48aec7ac350a03e62248e24c6a76b2/MaxText/max_utils.py#L739"""
    output_gb = bytes_to_gb(compiled_stats.output_size_in_bytes)
    temp_gb = bytes_to_gb(compiled_stats.temp_size_in_bytes)
    argument_gb = bytes_to_gb(compiled_stats.argument_size_in_bytes)
    alias_gb = bytes_to_gb(compiled_stats.alias_size_in_bytes)
    host_temp_gb = bytes_to_gb(compiled_stats.host_temp_size_in_bytes)
    total_gb = output_gb + temp_gb + argument_gb - alias_gb
    print(
        f"Total memory size: {total_gb:.1f} GB, Output size: {output_gb:.1f} GB, Temp size: {temp_gb:.1f} GB, "
        f"Argument size: {argument_gb:.1f} GB, Host temp size: {host_temp_gb:.1f} GB."
    )


def print_compiled_cost_analysis(cost_stats):
    flops = float(cost_stats.get("flops", 0.0))
    bytes_accessed = float(cost_stats.get("bytes accessed", 0.0))
    gb = bytes_to_gb(bytes_accessed) if bytes_accessed else 0.0
    intensity = (flops / bytes_accessed) if bytes_accessed else float("nan")
    print(
        f"FLOPs: {flops:.3e}, Bytes: {bytes_accessed:.3e} ({gb:.1f} GB), "
        f"Intensity: {intensity:.1f} FLOPs/byte"
    )


def print_mem_stats(label: str):
    """from: https://github.com/AI-Hypercomputer/maxtext/blob/7898576359bacde81be25cb3038e348aac1f943b/MaxText/max_utils.py#L713"""
    print(f"
Memstats: {label}:")
    try:
        for d in jax.local_devices():
            stats = d.memory_stats()
            used = round(stats["bytes_in_use"] / 2**30, 2)
            limit = round(stats["bytes_limit"] / 2**30, 2)
            print(f"	Using (GB) {used} / {limit} ({used/limit:%}) on {d}")
    except (RuntimeError, KeyError, TypeError) as ex:
        print(f"	Memstats unavailable, error: {ex}")

def _calculate_ffn_tflops_per_device(args):
  """Helper function to calculate matmul TFLOP in ffn based on MLP dimension.
  Adapted from https://github.com/AI-Hypercomputer/maxtext/blob/7898576359bacde81be25cb3038e348aac1f943b/MaxText/max_utils.py#L473

  Applies to:
    - Dense FFN layers (mlp_dim = config.mlp_dim).
    - MoE FFN layers (mlp_dim = config.moe_mlp_dim),
      need to scale by shared_experts or num_experts_per_tok.
  """

def calculate_tflops_per_device_tokenizer(args: TokenizerArgs):
    """Calculate training TFLOP for the Tokenizer."""
    input_dim = args.image_channels * (args.patch_size ** 2)

    tflops_per_device_encoder = calculate_tflops_per_device_transformer(args, input_dim, args.model_dim, args.ffn_dim, args.latent_dim, fully_causal=False)
    tflops_per_device_decoder = calculate_tflops_per_device_transformer(args, args.latent_dim, args.model_dim, args.ffn_dim, input_dim, fully_causal=False)

    batch_per_device = args.batch_size // jax.device_count()
    hn = (args.image_height + args.patch_size - 1) // args.patch_size
    wn = (args.image_width + args.patch_size - 1) // args.patch_size
    N = int(hn * wn)
    L = float(args.latent_dim)
    K = float(args.num_latents)

    # VQ distance compute (forward only; gradients do not flow through distances with STE)
    num_tokens = batch_per_device * args.seq_len * N
    vq_flops = 2.0 * num_tokens * L * K

    forward_flops = tflops_per_device_encoder + tflops_per_device_decoder

    # Training FLOPs: scale forward by ~3 for fwd+bwd+update, add VQ once
    total_training_flops = 3.0 * forward_flops + vq_flops

    to_tflops = 1.0 / 1e12
    result = {
        "encoder_tflops": tflops_per_device_encoder * to_tflops * 3.0,
        "decoder_tflops": tflops_per_device_decoder * to_tflops * 3.0,
        "vq_tflops": vq_flops * to_tflops * 3.0,
        "total_tflops": total_training_flops * to_tflops,
    }

    return result


def calculate_tflops_per_device_transformer(args: TokenizerArgs | LAMArgs | DynamicsArgs, input_dim: int, model_dim: int, ffn_dim: int, out_dim: int, fully_causal: bool = False):
    """Calculate training TFLOP for the STTransformer."""
    num_devices = jax.device_count()
    batch_per_device = args.batch_size // num_devices

    T = int(args.seq_len)
    H = int(args.image_height)
    W = int(args.image_width)

    # Number of patches per frame
    hn = (H + args.patch_size - 1) // args.patch_size
    wn = (W + args.patch_size - 1) // args.patch_size
    N = int(hn * wn)

    # Token counts for projections
    num_tokens = batch_per_device * T * N

    # Per-block FLOPs
    # Spatial attention (non-causal)
    qkv_spatial = 3.0 * 2.0 * num_tokens * model_dim * model_dim
    causal_multiplier = 0.5 if fully_causal else 1.0
    attn_spatial = causal_multiplier * 2.0 * 2.0 * (batch_per_device * T) * (N * N) * model_dim
    proj_spatial = 2.0 * num_tokens * model_dim * model_dim

    # Temporal attention (causal)
    qkv_temporal = 3.0 * 2.0 * num_tokens * model_dim * model_dim
    attn_temporal = 2.0 * (batch_per_device * N) * (T * T) * model_dim  # = 0.5 * (2 * 2 * B*N * T^2 * M)
    proj_temporal = 2.0 * num_tokens * model_dim * model_dim

    # FFN per block (2 matmuls)
    # FIXME (f.srambical):
    ffn_block = 4.0 * batch_per_device * T * N * model_dim * ffn_dim

    # NOTE (f.srambical): The calculation omits layer norm flops
    block_total = (
        qkv_spatial
        + attn_spatial
        + proj_spatial
        + qkv_temporal
        + attn_temporal
        + proj_temporal
        + ffn_block
    )

    # FIXME (f.srambical):
    input_output = 2.0 * num_tokens * (input_dim * model_dim + model_dim * out_dim)
    return block_total + input_output