From 086c9600a8017516e6e9e1f68222632fe0930fe1 Mon Sep 17 00:00:00 2001 From: Zihang-He Date: Sun, 25 May 2025 23:43:45 +0000 Subject: [PATCH] added gradient clipping --- fastvideo/v1/pipelines/training_pipeline.py | 16 +- fastvideo/v1/pipelines/training_utils.py | 203 ++++++++++++++++++++ train_v1.sh | 3 +- 3 files changed, 215 insertions(+), 7 deletions(-) diff --git a/fastvideo/v1/pipelines/training_pipeline.py b/fastvideo/v1/pipelines/training_pipeline.py index 8d019336d..9cbe9aa62 100644 --- a/fastvideo/v1/pipelines/training_pipeline.py +++ b/fastvideo/v1/pipelines/training_pipeline.py @@ -30,7 +30,7 @@ from fastvideo.v1.logger import init_logger from fastvideo.v1.pipelines import ComposedPipelineBase from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.v1.pipelines.wan.wan_pipeline import WanValidationPipeline - +from fastvideo.v1.pipelines.training_utils import _clip_grad_norm_while_handling_failing_dtensor_cases logger = init_logger(__name__) # Manual gradient checking flag - set to True to enable gradient verification @@ -683,17 +683,21 @@ class WanTrainingPipeline(TrainingPipeline): sp_group = get_sp_group() sp_group.all_reduce(avg_loss, op=torch.distributed.ReduceOp.AVG) total_loss += avg_loss.item() - - # TODO(will): clip grad norm - # grad_norm = transformer.clip_grad_norm_(max_grad_norm) + + model_parts = [self.transformer] + grad_norm = _clip_grad_norm_while_handling_failing_dtensor_cases( + [p for m in model_parts for p in m.parameters()], + max_grad_norm, + foreach=None, + ) + optimizer.step() print('device after optimizer step', next(transformer.named_parameters())[1].device) lr_scheduler.step() print('device after scheduler step', next(transformer.named_parameters())[1].device) - return total_loss, 0.0 - # return total_loss, grad_norm.item() + return total_loss, grad_norm.item() def forward( self, diff --git a/fastvideo/v1/pipelines/training_utils.py b/fastvideo/v1/pipelines/training_utils.py index e1a0dbe93..97ef5b73e 100644 --- a/fastvideo/v1/pipelines/training_utils.py +++ b/fastvideo/v1/pipelines/training_utils.py @@ -1,5 +1,10 @@ import json import os +from typing import List, Optional, Union, Tuple +import math +import torch +import torch.distributed.tensor +import torch.distributed as dist import torch from torch.distributed.fsdp import FullStateDictConfig @@ -8,6 +13,8 @@ from torch.distributed.fsdp import StateDictType from fastvideo.v1.logger import init_logger +_HAS_ERRORED_CLIP_GRAD_NORM_WHILE_HANDLING_FAILING_DTENSOR_CASES = False + logger = init_logger(__name__) @@ -39,3 +46,199 @@ def save_checkpoint(transformer, rank, output_dir, step): logger.info("--> checkpoint saved at step {step} to {weight_path}", step=step, weight_path=weight_path) + + +def _clip_grad_norm_while_handling_failing_dtensor_cases( + parameters: Union[torch.Tensor, List[torch.Tensor]], + max_norm: float, + norm_type: float = 2.0, + error_if_nonfinite: bool = False, + foreach: Optional[bool] = None, + pp_mesh: Optional[torch.distributed.device_mesh.DeviceMesh] = None, +) -> Optional[torch.Tensor]: + global _HAS_ERRORED_CLIP_GRAD_NORM_WHILE_HANDLING_FAILING_DTENSOR_CASES + + if not _HAS_ERRORED_CLIP_GRAD_NORM_WHILE_HANDLING_FAILING_DTENSOR_CASES: + try: + return clip_grad_norm_(parameters, max_norm, norm_type, error_if_nonfinite, foreach, pp_mesh) + except NotImplementedError as e: + if "DTensor does not support cross-mesh operation" in str(e): + # https://github.com/pytorch/pytorch/issues/134212 + logger.warning( + "DTensor does not support cross-mesh operation. If you haven't fully tensor-parallelized your " + "model, while combining other parallelisms such as FSDP, it could be the reason for this error. " + "Gradient clipping will be skipped and gradient norm will not be logged." + ) + except Exception as e: + logger.warning( + f"An error occurred while clipping gradients: {e}. Gradient clipping will be skipped and gradient " + f"norm will not be logged." + ) + _HAS_ERRORED_CLIP_GRAD_NORM_WHILE_HANDLING_FAILING_DTENSOR_CASES = True + return None + +# Copied from https://github.com/pytorch/torchtitan/blob/4a169701555ab9bd6ca3769f9650ae3386b84c6e/torchtitan/utils.py#L362 +@torch.no_grad() +def clip_grad_norm_( + parameters: Union[torch.Tensor, List[torch.Tensor]], + max_norm: float, + norm_type: float = 2.0, + error_if_nonfinite: bool = False, + foreach: Optional[bool] = None, + pp_mesh: Optional[torch.distributed.device_mesh.DeviceMesh] = None, +) -> torch.Tensor: + r""" + Clip the gradient norm of parameters. + + Gradient norm clipping requires computing the gradient norm over the entire model. + `torch.nn.utils.clip_grad_norm_` only computes gradient norm along DP/FSDP/TP dimensions. + We need to manually reduce the gradient norm across PP stages. + See https://github.com/pytorch/torchtitan/issues/596 for details. + + Args: + parameters (`torch.Tensor` or `List[torch.Tensor]`): + Tensors that will have gradients normalized. + max_norm (`float`): + Maximum norm of the gradients after clipping. + norm_type (`float`, defaults to `2.0`): + Type of p-norm to use. Can be `inf` for infinity norm. + error_if_nonfinite (`bool`, defaults to `False`): + If `True`, an error is thrown if the total norm of the gradients from `parameters` is `nan`, `inf`, or `-inf`. + foreach (`bool`, defaults to `None`): + Use the faster foreach-based implementation. If `None`, use the foreach implementation for CUDA and CPU native tensors + and silently fall back to the slow implementation for other device types. + pp_mesh (`torch.distributed.device_mesh.DeviceMesh`, defaults to `None`): + Pipeline parallel device mesh. If not `None`, will reduce gradient norm across PP stages. + + Returns: + `torch.Tensor`: + Total norm of the gradients + """ + grads = [p.grad for p in parameters if p.grad is not None] + + # TODO(aryan): Wait for next Pytorch release to use `torch.nn.utils.get_total_norm` + # total_norm = torch.nn.utils.get_total_norm(grads, norm_type, error_if_nonfinite, foreach) + total_norm = _get_total_norm(grads, norm_type, error_if_nonfinite, foreach) + + # If total_norm is a DTensor, the placements must be `torch.distributed._tensor.ops.math_ops._NormPartial`. + # We can simply reduce the DTensor to get the total norm in this tensor's process group + # and then convert it to a local tensor. + # It has two purposes: + # 1. to make sure the total norm is computed correctly when PP is used (see below) + # 2. to return a reduced total_norm tensor whose .item() would return the correct value + if isinstance(total_norm, torch.distributed.tensor.DTensor): + # Will reach here if any non-PP parallelism is used. + # If only using PP, total_norm will be a local tensor. + total_norm = total_norm.full_tensor() + + if pp_mesh is not None: + if math.isinf(norm_type): + dist.all_reduce(total_norm, op=dist.ReduceOp.MAX, group=pp_mesh.get_group()) + else: + total_norm **= norm_type + dist.all_reduce(total_norm, op=dist.ReduceOp.SUM, group=pp_mesh.get_group()) + total_norm **= 1.0 / norm_type + + _clip_grads_with_norm_(parameters, max_norm, total_norm, foreach) + return total_norm + + +@torch.no_grad() +def _clip_grads_with_norm_( + parameters: Union[torch.Tensor, List[torch.Tensor]], + max_norm: float, + total_norm: torch.Tensor, + foreach: Optional[bool] = None, +) -> None: + if isinstance(parameters, torch.Tensor): + parameters = [parameters] + grads = [p.grad for p in parameters if p.grad is not None] + max_norm = float(max_norm) + if len(grads) == 0: + return + grouped_grads: dict[Tuple[torch.device, torch.dtype], Tuple[List[List[torch.Tensor]], List[int]]] = ( + _group_tensors_by_device_and_dtype([grads]) + ) # type: ignore[assignment] + + clip_coef = max_norm / (total_norm + 1e-6) + + # Note: multiplying by the clamped coef is redundant when the coef is clamped to 1, but doing so + # avoids a `if clip_coef < 1:` conditional which can require a CPU <=> device synchronization + # when the gradients do not reside in CPU memory. + clip_coef_clamped = torch.clamp(clip_coef, max=1.0) + for (device, _), ([device_grads], _) in grouped_grads.items(): + if (foreach is None and _has_foreach_support(device_grads, device)) or ( + foreach and _device_has_foreach_support(device) + ): + torch._foreach_mul_(device_grads, clip_coef_clamped.to(device)) + elif foreach: + raise RuntimeError(f"foreach=True was passed, but can't use the foreach API on {device.type} tensors") + else: + clip_coef_clamped_device = clip_coef_clamped.to(device) + for g in device_grads: + g.mul_(clip_coef_clamped_device) + +def _get_total_norm( + tensors: Union[torch.Tensor, List[torch.Tensor]], + norm_type: float = 2.0, + error_if_nonfinite: bool = False, + foreach: Optional[bool] = None, +) -> torch.Tensor: + if isinstance(tensors, torch.Tensor): + tensors = [tensors] + else: + tensors = list(tensors) + norm_type = float(norm_type) + if len(tensors) == 0: + return torch.tensor(0.0) + first_device = tensors[0].device + grouped_tensors: dict[tuple[torch.device, torch.dtype], tuple[list[list[torch.Tensor]], list[int]]] = ( + _group_tensors_by_device_and_dtype( + [tensors] # type: ignore[list-item] + ) + ) # type: ignore[assignment] + + norms: List[torch.Tensor] = [] + for (device, _), ([device_tensors], _) in grouped_tensors.items(): + local_tensors = [ + t.to_local() if isinstance(t, torch.distributed.tensor.DTensor) else t + for t in device_tensors + ] + if (foreach is None and _has_foreach_support(local_tensors, device)) or ( + foreach and _device_has_foreach_support(device) + ): + norms.extend(torch._foreach_norm(local_tensors, norm_type)) + elif foreach: + raise RuntimeError(f"foreach=True was passed, but can't use the foreach API on {device.type} tensors") + else: + norms.extend([torch.linalg.vector_norm(g, norm_type) for g in local_tensors]) + + total_norm = torch.linalg.vector_norm(torch.stack([norm.to(first_device) for norm in norms]), norm_type) + + if error_if_nonfinite and torch.logical_or(total_norm.isnan(), total_norm.isinf()): + raise RuntimeError( + f"The total norm of order {norm_type} for gradients from " + "`parameters` is non-finite, so it cannot be clipped. To disable " + "this error and scale the gradients by the non-finite norm anyway, " + "set `error_if_nonfinite=False`" + ) + return total_norm + +def _get_foreach_kernels_supported_devices() -> list[str]: + r"""Return the device type list that supports foreach kernels.""" + return ["cuda", "xpu", torch._C._get_privateuse1_backend_name()] + +@torch.no_grad() +def _group_tensors_by_device_and_dtype( + tensorlistlist: List[List[Optional[torch.Tensor]]], + with_indices: bool = False, +) -> dict[tuple[torch.device, torch.dtype], tuple[List[List[Optional[torch.Tensor]]], List[int]]]: + return torch._C._group_tensors_by_device_and_dtype(tensorlistlist, with_indices) + + +def _device_has_foreach_support(device: torch.device) -> bool: + return device.type in (_get_foreach_kernels_supported_devices() + ["cpu"]) and not torch.jit.is_scripting() + + +def _has_foreach_support(tensors: List[torch.Tensor], device: torch.device) -> bool: + return _device_has_foreach_support(device) and all(t is None or type(t) in [torch.Tensor] for t in tensors) diff --git a/train_v1.sh b/train_v1.sh index 999cc9699..43b37c8a4 100644 --- a/train_v1.sh +++ b/train_v1.sh @@ -49,4 +49,5 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\ --multi_phased_distill_schedule "4000-1" \ --weight_decay 0.01 \ --not_apply_cfg_solver \ - --master_weight_type "bf16" \ No newline at end of file + --master_weight_type "bf16" \ + --max_grad_norm 1.0 \ \ No newline at end of file