Where should gradient clipping go in a mixed-precision LoRA training step?

I work on CloudGPU. When you use PyTorch gradient scaling, clip after scaler.unscale_(optimizer) and before scaler.step(optimizer). Clipping scaled gradients changes the meaning of the threshold. During gradient accumulation, wait until the effective batch is complete before unscaling.

1. Identify who owns the training step

The function below is for a custom, single-optimizer PyTorch loop using CUDA autocast with float16. It is not an instruction to add a second scaler inside a trainer that already manages mixed precision. In a managed trainer, use its documented clipping hook instead.

Construct the optimizer from the adapter parameters you actually intend to train. Frozen parameters do not become trainable just because you include them in a clipping call.

2. Put the operations in an explicit order

import torch


def adapter_step(model, optimizer, scaler, inputs, targets, loss_fn):
    # Inputs, targets and the prepared model already belong to CUDA.
    optimizer.zero_grad(set_to_none=True)
    trainable = [p for p in model.parameters() if p.requires_grad]
    if not trainable:
        raise ValueError("No trainable parameters")

    with torch.autocast(device_type="cuda", dtype=torch.float16):
        loss = loss_fn(model(inputs), targets)
    if loss.ndim != 0 or not torch.isfinite(loss.detach()).item():
        raise ValueError("Expected a finite scalar loss")

    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(
        trainable, max_norm=1.0, error_if_nonfinite=True,
    )
    scaler.step(optimizer)
    scaler.update()
    return loss.detach()

Create one torch.amp.GradScaler("cuda") for the training loop, outside this function, and pass it in. The clipping threshold is an illustrative setting, not a recommendation validated for your adapter. The function assumes one update per call and that the optimizer owns the intended trainable parameters.

The non-finite checks are a fail-fast diagnostic choice. A failed gradient check raises before the update; stop and investigate rather than catching it and silently continuing the same step. PyTorch's normal scaler-managed path can instead skip an update when gradients are non-finite. Choose the behavior deliberately.

3. Adjust the boundary for accumulation

If several microbatches form one update, zero once at the beginning of the accumulation group. Keep the scale unchanged while accumulating. Unscale and clip once after the final backward call, then step and update once. Weight each microbatch's loss according to the effective-batch reduction, including a shorter final group; do not paste this one-batch function unchanged into an accumulation loop.

Official PyTorch mixed-precision examples document unscaling, clipping and accumulation. Use your model's own forward/loss contract; the example assumes a tensor-returning model and a scalar loss.

For server-side task limits around a custom training job, see the CloudGPU integration guide.

Related CloudGPU video: This Video Was Made by AI. Now Run AI Yourself. It is remote-compute context, not a LoRA training demonstration.

This function was syntax-checked, not run as a training experiment. No convergence, memory or performance result is asserted.