nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

move utils into training_utils

Browse Source
This commit is contained in:
Will Lin
2025-05-28 11:40:15 -07:00
parent 086c9600a8
commit ec6ba01eb1
3 changed files with 113 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/v1/pipelines/training_pipeline.py
+6 -54
View File
@@ -1,5 +1,4 @@
import gc
import math
import os
import sys
import time
@@ -29,8 +28,11 @@ from fastvideo.v1.forward_context import set_forward_context
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.training_utils import (
_clip_grad_norm_while_handling_failing_dtensor_cases,
compute_density_for_timestep_sampling, get_sigmas)
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
@@ -38,56 +40,6 @@ ENABLE_GRADIENT_CHECK = False
GRADIENT_CHECK_DTYPE = torch.bfloat16
def compute_density_for_timestep_sampling(
weighting_scheme: str,
batch_size: int,
generator,
logit_mean: float = None,
logit_std: float = None,
mode_scale: float = None,
):
"""
Compute the density for sampling the timesteps when doing SD3 training.
Courtesy: This was contributed by Rafie Walker in https://github.com/huggingface/diffusers/pull/8528.
SD3 paper reference: https://arxiv.org/abs/2403.03206v1.
"""
if weighting_scheme == "logit_normal":
# See 3.1 in the SD3 paper ($rf/lognorm(0.00,1.00)$).
u = torch.normal(
mean=logit_mean,
std=logit_std,
size=(batch_size, ),
device="cpu",
generator=generator,
)
u = torch.nn.functional.sigmoid(u)
elif weighting_scheme == "mode":
u = torch.rand(size=(batch_size, ), device="cpu", generator=generator)
u = 1 - u - mode_scale * (torch.cos(math.pi * u / 2)**2 - 1 + u)
else:
u = torch.rand(size=(batch_size, ), device="cpu", generator=generator)
return u
def get_sigmas(noise_scheduler,
device,
timesteps,
n_dim=4,
dtype=torch.float32):
sigmas = noise_scheduler.sigmas.to(device=device, dtype=dtype)
schedule_timesteps = noise_scheduler.timesteps.to(device)
timesteps = timesteps.to(device)
step_indices = [(schedule_timesteps == t).nonzero().item()
for t in timesteps]
sigma = sigmas[step_indices].flatten()
while len(sigma.shape) < n_dim:
sigma = sigma.unsqueeze(-1)
return sigma
class TrainingPipeline(ComposedPipelineBase, ABC):
"""
A pipeline for training a model. All training pipelines should inherit from this class.
@@ -683,14 +635,14 @@ class WanTrainingPipeline(TrainingPipeline):
sp_group = get_sp_group()
sp_group.all_reduce(avg_loss, op=torch.distributed.ReduceOp.AVG)
total_loss += avg_loss.item()
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)
fastvideo/v1/pipelines/training_utils.py
+104 -34
View File
@@ -1,15 +1,16 @@
import json
import os
from typing import List, Optional, Union, Tuple
import math
from typing import List, Optional, Tuple, Union
import torch
import torch.distributed.tensor
import torch.distributed as dist
import torch
from torch.distributed.fsdp import FullStateDictConfig
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import StateDictType
import torch.distributed.tensor
from fastvideo.v1.logger import init_logger
@@ -18,6 +19,56 @@ _HAS_ERRORED_CLIP_GRAD_NORM_WHILE_HANDLING_FAILING_DTENSOR_CASES = False
logger = init_logger(__name__)
def compute_density_for_timestep_sampling(
weighting_scheme: str,
batch_size: int,
generator,
logit_mean: float = None,
logit_std: float = None,
mode_scale: float = None,
):
"""
Compute the density for sampling the timesteps when doing SD3 training.
Courtesy: This was contributed by Rafie Walker in https://github.com/huggingface/diffusers/pull/8528.
SD3 paper reference: https://arxiv.org/abs/2403.03206v1.
"""
if weighting_scheme == "logit_normal":
# See 3.1 in the SD3 paper ($rf/lognorm(0.00,1.00)$).
u = torch.normal(
mean=logit_mean,
std=logit_std,
size=(batch_size, ),
device="cpu",
generator=generator,
)
u = torch.nn.functional.sigmoid(u)
elif weighting_scheme == "mode":
u = torch.rand(size=(batch_size, ), device="cpu", generator=generator)
u = 1 - u - mode_scale * (torch.cos(math.pi * u / 2)**2 - 1 + u)
else:
u = torch.rand(size=(batch_size, ), device="cpu", generator=generator)
return u
def get_sigmas(noise_scheduler,
device,
timesteps,
n_dim=4,
dtype=torch.float32):
sigmas = noise_scheduler.sigmas.to(device=device, dtype=dtype)
schedule_timesteps = noise_scheduler.timesteps.to(device)
timesteps = timesteps.to(device)
step_indices = [(schedule_timesteps == t).nonzero().item()
for t in timesteps]
sigma = sigmas[step_indices].flatten()
while len(sigma.shape) < n_dim:
sigma = sigma.unsqueeze(-1)
return sigma
def save_checkpoint(transformer, rank, output_dir, step):
# Configure FSDP to save full state dict
FSDP.set_state_dict_type(
@@ -60,7 +111,8 @@ def _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)
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
@@ -72,11 +124,11 @@ def _clip_grad_norm_while_handling_failing_dtensor_cases(
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."
)
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_(
@@ -133,10 +185,14 @@ def clip_grad_norm_(
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())
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())
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)
@@ -156,10 +212,11 @@ def _clip_grads_with_norm_(
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]
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
@@ -168,16 +225,18 @@ def _clip_grads_with_norm_(
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)
):
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")
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,
@@ -192,53 +251,64 @@ def _get_total_norm(
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]
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
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)
):
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")
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])
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)
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()):
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`"
)
"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)
) -> 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()
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)
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)
train_v1.sh
+3 -3
View File
@@ -30,8 +30,8 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\
--max_train_steps=120 \
--learning_rate=1e-6\
--mixed_precision="bf16"\
--checkpointing_steps=10 \
--validation_steps 5\
--checkpointing_steps=60 \
--validation_steps 20\
--validation_sampling_steps "2,4,8" \
--log_validation \
--checkpoints_total_limit 3\
@@ -50,4 +50,4 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\
--weight_decay 0.01 \
--not_apply_cfg_solver \
--master_weight_type "bf16" \
--max_grad_norm 1.0 \
--max_grad_norm 1.0