File size: 3,156 Bytes
b9fc164 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | """Shared loss and optimizer handling for distributed training."""
import torch
import torch.distributed as dist
from torch.distributed.tensor import DTensor
from torch.distributed.fsdp import FSDPModule
from src.distributed.communications import all_ranks_finite
def build_token_timesteps(timesteps: torch.Tensor, latents: torch.Tensor, model) -> torch.Tensor:
"""Repeat each patch-frame timestep across its spatial tokens."""
patch_time, patch_height, patch_width = model.config.patch_size
frames, height, width = latents.shape[-3:]
spatial_tokens = (height // patch_height) * (width // patch_width)
return timesteps[:, ::patch_time].unsqueeze(-1).expand(-1, -1, spatial_tokens).reshape(
timesteps.shape[0], (frames // patch_time) * spatial_tokens,
)
def synchronize_batch_frames(batch: dict, device: torch.device) -> dict:
"""Keep chunk-by-chunk forward counts equal when ranks draw different buckets."""
if not dist.is_initialized() or dist.get_world_size() == 1:
return batch
frames = torch.tensor(batch["pixel"].shape[3], dtype=torch.int64, device=device)
dist.all_reduce(frames, op=dist.ReduceOp.MIN)
return {**batch, "pixel": batch["pixel"][:, :, :, :int(frames.item())]}
def backward_if_finite(loss: torch.Tensor, logger, context: str, *, model=None) -> bool:
if not all_ranks_finite(loss):
# A skipped backward leaves FSDP's root parameters unsharded. Restore
# the optimizer's sharded parameter views without dropping prior grads.
if model is not None:
for module in reversed(list(model.modules())):
if isinstance(module, FSDPModule):
module.reshard()
logger.warning(f"{context}: non-finite loss on at least one rank; skipping backward on all ranks.")
return False
loss.backward()
return True
def step_optimizer(optimizer, scheduler, max_grad_norm, successful_backwards: int, logger):
"""Skip empty/non-finite updates and advance the scheduler only after a step."""
if successful_backwards == 0:
optimizer.zero_grad(set_to_none=True)
return None
grad_norm = torch.nn.utils.clip_grad_norm_(
[parameter for group in optimizer.param_groups for parameter in group["params"]],
max_norm=float("inf") if max_grad_norm is None else max_grad_norm,
)
if isinstance(grad_norm, DTensor):
grad_norm = grad_norm.full_tensor()
if all_ranks_finite(grad_norm):
optimizer.step()
scheduler.step()
else:
logger.warning("Non-finite gradient norm on at least one rank; skipping optimizer and scheduler on all ranks.")
optimizer.zero_grad(set_to_none=True)
return grad_norm
def log_bucket_batch_sizes(cfg, world_size: int, logger) -> None:
buckets = list(cfg.bucket_configs or [])
sizes = sorted({bucket[0] for bucket in buckets})
logger.info(f"Per-rank bucket batch sizes: {sizes}")
logger.info(
"Effective batch size per update = sum of rank-local bucket batch sizes "
f"across {world_size} ranks and {cfg.gradient_accumulation_steps} accumulation steps."
)
|