"""Execution-only optimizations; model tensors and inference contract are unchanged."""
import torch
import torch.distributed as dist

def install_fast_execution(model, precision):
    assert precision in ('fp32', 'tf32')
    torch.backends.cuda.matmul.allow_tf32 = precision == 'tf32'
    model._mot_driver.mot_checkpoint_mixed_attn = False
    model.set_training_runtime(use_gradient_checkpointing=False,
        use_gradient_checkpointing_offload=False, max_timestep_boundary=1., min_timestep_boundary=0.)

def sum_gradients_fast(parameters, world):
    """SUM global-accumulation-scaled FP32 gradients in one collective."""
    assert parameters and all(p.dtype == torch.float32 for p in parameters)
    flat = torch.empty(sum(p.numel() for p in parameters), device=parameters[0].device, dtype=torch.float32)
    views = []
    offset = 0
    for p in parameters:
        views.append(flat[offset:offset+p.numel()].view_as(p))
        offset += p.numel()
    def reduce():
        for p, view in zip(parameters, views):
            if p.grad is None: view.zero_()
            else: view.copy_(p.grad)
        if world > 1: dist.all_reduce(flat, op=dist.ReduceOp.SUM)
        for p, view in zip(parameters, views): p.grad = view
    return reduce