[Refactor] Switch to FSDP (#42)
This commit is contained in:
Zhang Peiyuan
+3
-1
@@ -48,4 +48,6 @@ samples/
|
||||
*validation/
|
||||
data/
|
||||
outputs/
|
||||
outputs_video
|
||||
outputs_video
|
||||
sbatch.sh
|
||||
*.out
|
||||
@@ -23,9 +23,10 @@ We've prepared some debug data to facilitate development. To make sure the train
|
||||
|
||||
```
|
||||
python scripts/download_hf.py --repo_id=Stealths-Video/dummyVid --local_dir=data/dummyVid --repo_type=dataset
|
||||
python scripts/download_hf.py --repo_id=Stealths-Video/mochi --local_dir=data/mochi --repo_type=model
|
||||
python scripts/download_hf.py --repo_id=Stealths-Video/mochi_diffuser --local_dir=data/mochi --repo_type=model
|
||||
python scripts/download_hf.py --repo_id=Stealths-Video/Mochi-Synthetic-Data --local_dir=data/Mochi-Synthetic-Data --repo_type=dataset
|
||||
python scripts/download_hf.py --repo_id=Stealths-Video/Encoder_Overfit_Data --local_dir=data/Encoder_Overfit_Data --repo_type=dataset
|
||||
python scripts/download_hf.py --repo_id=Stealths-Video/BW-Finetune-Dataset --local_dir=data/BW-Finetune-Dataset --repo_type=dataset
|
||||
```
|
||||
|
||||
## How to overfit
|
||||
|
||||
@@ -22,6 +22,7 @@ class LatentDataset(Dataset):
|
||||
self.prompt_attention_mask_dir = os.path.join(self.datase_dir_path, "prompt_attention_mask")
|
||||
with open(self.json_path, 'r') as f:
|
||||
self.data_anno = json.load(f)
|
||||
self.data_anno = sorted(self.data_anno, key=lambda x: x['latent_path'])
|
||||
self.num_latent_t = num_latent_t
|
||||
self.uncond_prompt_embed = torch.load(os.path.join(uncond_prompt_embed_mask_dir, "embed.pt"), map_location="cpu", weights_only=True)
|
||||
self.uncond_prompt_mask = torch.load(os.path.join(uncond_prompt_embed_mask_dir, "mask.pt"), map_location="cpu", weights_only=True)
|
||||
|
||||
@@ -0,0 +1,86 @@
|
||||
import torch
|
||||
import os
|
||||
import torch.distributed as dist
|
||||
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
|
||||
checkpoint_wrapper,
|
||||
CheckpointImpl,
|
||||
apply_activation_checkpointing,
|
||||
)
|
||||
|
||||
from torch.distributed.fsdp import (
|
||||
FullyShardedDataParallel as FSDP,
|
||||
StateDictType,
|
||||
FullStateDictConfig, # general model non-sharded, non-flattened params
|
||||
LocalStateDictConfig, # flattened params, usable only by FSDP
|
||||
# ShardedStateDictConfig, # un-flattened param but shards, usable by other parallel schemes.
|
||||
)
|
||||
|
||||
from fastvideo.model.modeling_mochi import MochiTransformerBlock
|
||||
|
||||
from functools import partial
|
||||
|
||||
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
|
||||
|
||||
from torch.distributed.fsdp import MixedPrecision, ShardingStrategy
|
||||
import functools
|
||||
|
||||
|
||||
non_reentrant_wrapper = partial(
|
||||
checkpoint_wrapper,
|
||||
checkpoint_impl=CheckpointImpl.NO_REENTRANT,
|
||||
)
|
||||
|
||||
check_fn = lambda submodule: isinstance(submodule, MochiTransformerBlock)
|
||||
|
||||
|
||||
def apply_fsdp_checkpointing(model):
|
||||
"""apply activation checkpointing to model
|
||||
returns None as model is updated directly
|
||||
"""
|
||||
print(f"--> applying fdsp activation checkpointing...")
|
||||
|
||||
apply_activation_checkpointing(
|
||||
model, checkpoint_wrapper_fn=non_reentrant_wrapper, check_fn=check_fn
|
||||
)
|
||||
|
||||
bf16_mix = MixedPrecision(
|
||||
param_dtype=torch.bfloat16,
|
||||
# Gradient communication precision.
|
||||
reduce_dtype=torch.bfloat16,
|
||||
# Buffer precision.
|
||||
buffer_dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
|
||||
|
||||
def get_fsdp_kwargs( sharding_strategy, cpu_offload=False):
|
||||
auto_wrap_policy = functools.partial(
|
||||
transformer_auto_wrap_policy,
|
||||
transformer_layer_cls={
|
||||
MochiTransformerBlock,
|
||||
},
|
||||
)
|
||||
|
||||
mixed_precision = bf16_mix
|
||||
|
||||
if sharding_strategy == "full":
|
||||
sharding_strategy = ShardingStrategy.FULL_SHARD
|
||||
elif sharding_strategy == "none":
|
||||
sharding_strategy = ShardingStrategy.NO_SHARD
|
||||
auto_wrap_policy = None
|
||||
elif sharding_strategy == "hybrid_zero2":
|
||||
sharding_strategy = ShardingStrategy._HYBRID_SHARD_ZERO2
|
||||
|
||||
device_id = torch.cuda.current_device()
|
||||
cpu_offload=torch.distributed.fsdp.CPUOffload(offload_params=True) if cpu_offload else None
|
||||
return {
|
||||
"auto_wrap_policy": auto_wrap_policy,
|
||||
"mixed_precision": mixed_precision,
|
||||
"sharding_strategy": sharding_strategy,
|
||||
"device_id": device_id,
|
||||
"limit_all_gathers": True,
|
||||
"cpu_offload": cpu_offload,
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -1,161 +0,0 @@
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
import diffusers
|
||||
from diffusers.models.attention_processor import Attention
|
||||
from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
|
||||
from fastvideo.utils.communications import all_gather, all_to_all_4D
|
||||
|
||||
class NewMochiAttnProcessor2_0:
|
||||
"""Attention processor used in Mochi."""
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn: Attention,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
image_rotary_emb: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
# [b, s, h * d]
|
||||
query = attn.to_q(hidden_states)
|
||||
key = attn.to_k(hidden_states)
|
||||
value = attn.to_v(hidden_states)
|
||||
|
||||
# [b, s, h=24, d=128]
|
||||
query = query.unflatten(2, (attn.heads, -1))
|
||||
key = key.unflatten(2, (attn.heads, -1))
|
||||
value = value.unflatten(2, (attn.heads, -1))
|
||||
|
||||
|
||||
if attn.norm_q is not None:
|
||||
query = attn.norm_q(query)
|
||||
if attn.norm_k is not None:
|
||||
key = attn.norm_k(key)
|
||||
# [b, 256, h * d]
|
||||
encoder_query = attn.add_q_proj(encoder_hidden_states)
|
||||
encoder_key = attn.add_k_proj(encoder_hidden_states)
|
||||
encoder_value = attn.add_v_proj(encoder_hidden_states)
|
||||
|
||||
# [b, 256, h=24, d=128]
|
||||
encoder_query = encoder_query.unflatten(2, (attn.heads, -1))
|
||||
encoder_key = encoder_key.unflatten(2, (attn.heads, -1))
|
||||
encoder_value = encoder_value.unflatten(2, (attn.heads, -1))
|
||||
|
||||
|
||||
if attn.norm_added_q is not None:
|
||||
encoder_query = attn.norm_added_q(encoder_query)
|
||||
if attn.norm_added_k is not None:
|
||||
encoder_key = attn.norm_added_k(encoder_key)
|
||||
|
||||
if image_rotary_emb is not None:
|
||||
freqs_cos, freqs_sin = image_rotary_emb[0], image_rotary_emb[1]
|
||||
# shard the head dimension
|
||||
if get_sequence_parallel_state():
|
||||
# B, S, H, D to (S, B,) H, D
|
||||
# batch_size, seq_len, attn_heads, head_dim
|
||||
query = all_to_all_4D(query, scatter_dim=2, gather_dim=1)
|
||||
key = all_to_all_4D(key, scatter_dim=2, gather_dim=1)
|
||||
value = all_to_all_4D(value, scatter_dim=2, gather_dim=1)
|
||||
|
||||
|
||||
def shrink_head(encoder_state, dim):
|
||||
local_heads = encoder_state.shape[dim] // nccl_info.sp_size
|
||||
return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads)
|
||||
encoder_query = shrink_head(encoder_query, dim=2)
|
||||
encoder_key = shrink_head(encoder_key, dim=2)
|
||||
encoder_value = shrink_head(encoder_value, dim=2)
|
||||
if image_rotary_emb is not None:
|
||||
freqs_cos = shrink_head(freqs_cos, dim=1)
|
||||
freqs_sin = shrink_head(freqs_sin, dim=1)
|
||||
|
||||
|
||||
|
||||
if image_rotary_emb is not None:
|
||||
def apply_rotary_emb(x, freqs_cos, freqs_sin):
|
||||
x_even = x[..., 0::2].float()
|
||||
x_odd = x[..., 1::2].float()
|
||||
cos = (x_even * freqs_cos - x_odd * freqs_sin).to(x.dtype)
|
||||
sin = (x_even * freqs_sin + x_odd * freqs_cos).to(x.dtype)
|
||||
|
||||
return torch.stack([cos, sin], dim=-1).flatten(-2)
|
||||
query = apply_rotary_emb(query, freqs_cos, freqs_sin)
|
||||
key = apply_rotary_emb(key, freqs_cos, freqs_sin)
|
||||
|
||||
query, key, value = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2)
|
||||
encoder_query, encoder_key, encoder_value = (
|
||||
encoder_query.transpose(1, 2),
|
||||
encoder_key.transpose(1, 2),
|
||||
encoder_value.transpose(1, 2),
|
||||
)
|
||||
# [b, h, s, d]
|
||||
sequence_length = query.size(2)
|
||||
encoder_sequence_length = encoder_query.size(2)
|
||||
|
||||
# Hint: please check encoder_query.shape
|
||||
query = torch.cat([query, encoder_query], dim=2)
|
||||
key = torch.cat([key, encoder_key], dim=2)
|
||||
value = torch.cat([value, encoder_value], dim=2)
|
||||
|
||||
|
||||
|
||||
hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
|
||||
|
||||
if get_sequence_parallel_state():
|
||||
hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
|
||||
(sequence_length, encoder_sequence_length), dim=2
|
||||
)
|
||||
# B, H, S, D
|
||||
hidden_states = all_to_all_4D(hidden_states, scatter_dim=2, gather_dim=1)
|
||||
encoder_hidden_states = all_gather(encoder_hidden_states, dim=1).contiguous()
|
||||
hidden_states = hidden_states.transpose(1,2).flatten(2, 3)
|
||||
hidden_states = hidden_states.to(query.dtype)
|
||||
encoder_hidden_states = encoder_hidden_states.transpose(1, 2).flatten(2, 3)
|
||||
encoder_hidden_states = encoder_hidden_states.to(query.dtype)
|
||||
else:
|
||||
hidden_states = hidden_states.transpose(1, 2).flatten(2, 3)
|
||||
hidden_states = hidden_states.to(query.dtype)
|
||||
|
||||
hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
|
||||
(sequence_length, encoder_sequence_length), dim=1
|
||||
)
|
||||
|
||||
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
if hasattr(attn, "to_add_out"):
|
||||
encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
|
||||
|
||||
return hidden_states, encoder_hidden_states
|
||||
|
||||
|
||||
|
||||
|
||||
class NewMochiRoPE(nn.Module):
|
||||
def _get_positions(
|
||||
self,
|
||||
num_frames: int,
|
||||
height: int,
|
||||
width: int,
|
||||
device: Optional[torch.device] = None,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
) -> torch.Tensor:
|
||||
scale = (self.target_area / (height * width)) ** 0.5
|
||||
t = torch.arange(num_frames * nccl_info.sp_size, device=device, dtype=dtype)
|
||||
h = self._centers(-height * scale / 2, height * scale / 2, height, device, dtype)
|
||||
w = self._centers(-width * scale / 2, width * scale / 2, width, device, dtype)
|
||||
|
||||
grid_t, grid_h, grid_w = torch.meshgrid(t, h, w, indexing="ij")
|
||||
|
||||
positions = torch.stack([grid_t, grid_h, grid_w], dim=-1).view(-1, 3)
|
||||
return positions
|
||||
|
||||
def hf_mochi_add_sp_monkey_patch():
|
||||
diffusers.models.attention_processor.MochiAttnProcessor2_0.__call__ = NewMochiAttnProcessor2_0.__call__
|
||||
diffusers.models.transformers.transformer_mochi.MochiRoPE._get_positions = NewMochiRoPE._get_positions
|
||||
@@ -0,0 +1,527 @@
|
||||
# Copyright 2024 The Genmo team and The HuggingFace Team.
|
||||
# All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.utils import is_torch_version, logging
|
||||
from diffusers.utils.torch_utils import maybe_allow_in_graph
|
||||
from diffusers.models.attention import FeedForward
|
||||
from diffusers.models.attention_processor import Attention
|
||||
from diffusers.models.embeddings import MochiCombinedTimestepCaptionEmbedding, PatchEmbed
|
||||
from diffusers.models.modeling_outputs import Transformer2DModelOutput
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
from diffusers.models.normalization import AdaLayerNormContinuous, LuminaLayerNormContinuous, MochiRMSNormZero, RMSNorm
|
||||
from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
|
||||
from fastvideo.utils.communications import all_gather, all_to_all_4D
|
||||
import torch.nn.functional as F
|
||||
from diffusers.utils.torch_utils import is_torch_version, maybe_allow_in_graph
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
class MochiAttnProcessor2_0:
|
||||
"""Attention processor used in Mochi."""
|
||||
|
||||
def __init__(self):
|
||||
if not hasattr(F, "scaled_dot_product_attention"):
|
||||
raise ImportError("MochiAttnProcessor2_0 requires PyTorch 2.0. To use it, please upgrade PyTorch to 2.0.")
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn: Attention,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
image_rotary_emb: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
# [b, s, h * d]
|
||||
query = attn.to_q(hidden_states)
|
||||
key = attn.to_k(hidden_states)
|
||||
value = attn.to_v(hidden_states)
|
||||
|
||||
# [b, s, h=24, d=128]
|
||||
query = query.unflatten(2, (attn.heads, -1))
|
||||
key = key.unflatten(2, (attn.heads, -1))
|
||||
value = value.unflatten(2, (attn.heads, -1))
|
||||
|
||||
|
||||
if attn.norm_q is not None:
|
||||
query = attn.norm_q(query)
|
||||
if attn.norm_k is not None:
|
||||
key = attn.norm_k(key)
|
||||
# [b, 256, h * d]
|
||||
encoder_query = attn.add_q_proj(encoder_hidden_states)
|
||||
encoder_key = attn.add_k_proj(encoder_hidden_states)
|
||||
encoder_value = attn.add_v_proj(encoder_hidden_states)
|
||||
|
||||
# [b, 256, h=24, d=128]
|
||||
encoder_query = encoder_query.unflatten(2, (attn.heads, -1))
|
||||
encoder_key = encoder_key.unflatten(2, (attn.heads, -1))
|
||||
encoder_value = encoder_value.unflatten(2, (attn.heads, -1))
|
||||
|
||||
|
||||
if attn.norm_added_q is not None:
|
||||
encoder_query = attn.norm_added_q(encoder_query)
|
||||
if attn.norm_added_k is not None:
|
||||
encoder_key = attn.norm_added_k(encoder_key)
|
||||
|
||||
if image_rotary_emb is not None:
|
||||
freqs_cos, freqs_sin = image_rotary_emb[0], image_rotary_emb[1]
|
||||
# shard the head dimension
|
||||
if get_sequence_parallel_state():
|
||||
# B, S, H, D to (S, B,) H, D
|
||||
# batch_size, seq_len, attn_heads, head_dim
|
||||
query = all_to_all_4D(query, scatter_dim=2, gather_dim=1)
|
||||
key = all_to_all_4D(key, scatter_dim=2, gather_dim=1)
|
||||
value = all_to_all_4D(value, scatter_dim=2, gather_dim=1)
|
||||
|
||||
|
||||
def shrink_head(encoder_state, dim):
|
||||
local_heads = encoder_state.shape[dim] // nccl_info.sp_size
|
||||
return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads)
|
||||
encoder_query = shrink_head(encoder_query, dim=2)
|
||||
encoder_key = shrink_head(encoder_key, dim=2)
|
||||
encoder_value = shrink_head(encoder_value, dim=2)
|
||||
if image_rotary_emb is not None:
|
||||
freqs_cos = shrink_head(freqs_cos, dim=1)
|
||||
freqs_sin = shrink_head(freqs_sin, dim=1)
|
||||
|
||||
|
||||
|
||||
if image_rotary_emb is not None:
|
||||
def apply_rotary_emb(x, freqs_cos, freqs_sin):
|
||||
x_even = x[..., 0::2].float()
|
||||
x_odd = x[..., 1::2].float()
|
||||
cos = (x_even * freqs_cos - x_odd * freqs_sin).to(x.dtype)
|
||||
sin = (x_even * freqs_sin + x_odd * freqs_cos).to(x.dtype)
|
||||
|
||||
return torch.stack([cos, sin], dim=-1).flatten(-2)
|
||||
query = apply_rotary_emb(query, freqs_cos, freqs_sin)
|
||||
key = apply_rotary_emb(key, freqs_cos, freqs_sin)
|
||||
|
||||
query, key, value = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2)
|
||||
encoder_query, encoder_key, encoder_value = (
|
||||
encoder_query.transpose(1, 2),
|
||||
encoder_key.transpose(1, 2),
|
||||
encoder_value.transpose(1, 2),
|
||||
)
|
||||
# [b, h, s, d]
|
||||
sequence_length = query.size(2)
|
||||
encoder_sequence_length = encoder_query.size(2)
|
||||
|
||||
# Hint: please check encoder_query.shape
|
||||
query = torch.cat([query, encoder_query], dim=2)
|
||||
key = torch.cat([key, encoder_key], dim=2)
|
||||
value = torch.cat([value, encoder_value], dim=2)
|
||||
|
||||
|
||||
|
||||
hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
|
||||
|
||||
if get_sequence_parallel_state():
|
||||
hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
|
||||
(sequence_length, encoder_sequence_length), dim=2
|
||||
)
|
||||
# B, H, S, D
|
||||
hidden_states = all_to_all_4D(hidden_states, scatter_dim=2, gather_dim=1)
|
||||
encoder_hidden_states = all_gather(encoder_hidden_states, dim=1).contiguous()
|
||||
hidden_states = hidden_states.transpose(1,2).flatten(2, 3)
|
||||
hidden_states = hidden_states.to(query.dtype)
|
||||
encoder_hidden_states = encoder_hidden_states.transpose(1, 2).flatten(2, 3)
|
||||
encoder_hidden_states = encoder_hidden_states.to(query.dtype)
|
||||
else:
|
||||
hidden_states = hidden_states.transpose(1, 2).flatten(2, 3)
|
||||
hidden_states = hidden_states.to(query.dtype)
|
||||
|
||||
hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
|
||||
(sequence_length, encoder_sequence_length), dim=1
|
||||
)
|
||||
|
||||
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
if hasattr(attn, "to_add_out"):
|
||||
encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
|
||||
|
||||
return hidden_states, encoder_hidden_states
|
||||
|
||||
@maybe_allow_in_graph
|
||||
class MochiTransformerBlock(nn.Module):
|
||||
r"""
|
||||
Transformer block used in [Mochi](https://huggingface.co/genmo/mochi-1-preview).
|
||||
|
||||
Args:
|
||||
dim (`int`):
|
||||
The number of channels in the input and output.
|
||||
num_attention_heads (`int`):
|
||||
The number of heads to use for multi-head attention.
|
||||
attention_head_dim (`int`):
|
||||
The number of channels in each head.
|
||||
qk_norm (`str`, defaults to `"rms_norm"`):
|
||||
The normalization layer to use.
|
||||
activation_fn (`str`, defaults to `"swiglu"`):
|
||||
Activation function to use in feed-forward.
|
||||
context_pre_only (`bool`, defaults to `False`):
|
||||
Whether or not to process context-related conditions with additional layers.
|
||||
eps (`float`, defaults to `1e-6`):
|
||||
Epsilon value for normalization layers.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
num_attention_heads: int,
|
||||
attention_head_dim: int,
|
||||
pooled_projection_dim: int,
|
||||
qk_norm: str = "rms_norm",
|
||||
activation_fn: str = "swiglu",
|
||||
context_pre_only: bool = False,
|
||||
eps: float = 1e-6,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.context_pre_only = context_pre_only
|
||||
self.ff_inner_dim = (4 * dim * 2) // 3
|
||||
self.ff_context_inner_dim = (4 * pooled_projection_dim * 2) // 3
|
||||
|
||||
self.norm1 = MochiRMSNormZero(dim, 4 * dim, eps=eps, elementwise_affine=False)
|
||||
|
||||
if not context_pre_only:
|
||||
self.norm1_context = MochiRMSNormZero(dim, 4 * pooled_projection_dim, eps=eps, elementwise_affine=False)
|
||||
else:
|
||||
self.norm1_context = LuminaLayerNormContinuous(
|
||||
embedding_dim=pooled_projection_dim,
|
||||
conditioning_embedding_dim=dim,
|
||||
eps=eps,
|
||||
elementwise_affine=False,
|
||||
norm_type="rms_norm",
|
||||
out_dim=None,
|
||||
)
|
||||
|
||||
self.attn1 = Attention(
|
||||
query_dim=dim,
|
||||
cross_attention_dim=None,
|
||||
heads=num_attention_heads,
|
||||
dim_head=attention_head_dim,
|
||||
bias=False,
|
||||
qk_norm=qk_norm,
|
||||
added_kv_proj_dim=pooled_projection_dim,
|
||||
added_proj_bias=False,
|
||||
out_dim=dim,
|
||||
out_context_dim=pooled_projection_dim,
|
||||
context_pre_only=context_pre_only,
|
||||
processor=MochiAttnProcessor2_0(),
|
||||
eps=eps,
|
||||
elementwise_affine=True,
|
||||
)
|
||||
|
||||
# TODO(aryan): norm_context layers are not needed when `context_pre_only` is True
|
||||
self.norm2 = RMSNorm(dim, eps=eps, elementwise_affine=False)
|
||||
self.norm2_context = RMSNorm(pooled_projection_dim, eps=eps, elementwise_affine=False)
|
||||
|
||||
self.norm3 = RMSNorm(dim, eps=eps, elementwise_affine=False)
|
||||
self.norm3_context = RMSNorm(pooled_projection_dim, eps=eps, elementwise_affine=False)
|
||||
|
||||
self.ff = FeedForward(dim, inner_dim=self.ff_inner_dim, activation_fn=activation_fn, bias=False)
|
||||
self.ff_context = None
|
||||
if not context_pre_only:
|
||||
self.ff_context = FeedForward(
|
||||
pooled_projection_dim,
|
||||
inner_dim=self.ff_context_inner_dim,
|
||||
activation_fn=activation_fn,
|
||||
bias=False,
|
||||
)
|
||||
|
||||
self.norm4 = RMSNorm(dim, eps=eps, elementwise_affine=False)
|
||||
self.norm4_context = RMSNorm(pooled_projection_dim, eps=eps, elementwise_affine=False)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
temb: torch.Tensor,
|
||||
image_rotary_emb: Optional[torch.Tensor] = None,
|
||||
output_attn = False,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(hidden_states, temb)
|
||||
|
||||
if not self.context_pre_only:
|
||||
norm_encoder_hidden_states, enc_gate_msa, enc_scale_mlp, enc_gate_mlp = self.norm1_context(
|
||||
encoder_hidden_states, temb
|
||||
)
|
||||
else:
|
||||
norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states, temb)
|
||||
|
||||
attn_hidden_states, context_attn_hidden_states = self.attn1(
|
||||
hidden_states=norm_hidden_states,
|
||||
encoder_hidden_states=norm_encoder_hidden_states,
|
||||
image_rotary_emb=image_rotary_emb,
|
||||
)
|
||||
|
||||
hidden_states = hidden_states + self.norm2(attn_hidden_states) * torch.tanh(gate_msa).unsqueeze(1)
|
||||
norm_hidden_states = self.norm3(hidden_states) * (1 + scale_mlp.unsqueeze(1))
|
||||
ff_output = self.ff(norm_hidden_states)
|
||||
hidden_states = hidden_states + self.norm4(ff_output) * torch.tanh(gate_mlp).unsqueeze(1)
|
||||
|
||||
if not self.context_pre_only:
|
||||
encoder_hidden_states = encoder_hidden_states + self.norm2_context(
|
||||
context_attn_hidden_states
|
||||
) * torch.tanh(enc_gate_msa).unsqueeze(1)
|
||||
norm_encoder_hidden_states = self.norm3_context(encoder_hidden_states) * (1 + enc_scale_mlp.unsqueeze(1))
|
||||
context_ff_output = self.ff_context(norm_encoder_hidden_states)
|
||||
encoder_hidden_states = encoder_hidden_states + self.norm4_context(context_ff_output) * torch.tanh(
|
||||
enc_gate_mlp
|
||||
).unsqueeze(1)
|
||||
|
||||
if not output_attn:
|
||||
attn_hidden_states = None
|
||||
return hidden_states, encoder_hidden_states, attn_hidden_states
|
||||
|
||||
|
||||
class MochiRoPE(nn.Module):
|
||||
r"""
|
||||
RoPE implementation used in [Mochi](https://huggingface.co/genmo/mochi-1-preview).
|
||||
|
||||
Args:
|
||||
base_height (`int`, defaults to `192`):
|
||||
Base height used to compute interpolation scale for rotary positional embeddings.
|
||||
base_width (`int`, defaults to `192`):
|
||||
Base width used to compute interpolation scale for rotary positional embeddings.
|
||||
"""
|
||||
|
||||
def __init__(self, base_height: int = 192, base_width: int = 192) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.target_area = base_height * base_width
|
||||
|
||||
def _centers(self, start, stop, num, device, dtype) -> torch.Tensor:
|
||||
edges = torch.linspace(start, stop, num + 1, device=device, dtype=dtype)
|
||||
return (edges[:-1] + edges[1:]) / 2
|
||||
|
||||
def _get_positions(
|
||||
self,
|
||||
num_frames: int,
|
||||
height: int,
|
||||
width: int,
|
||||
device: Optional[torch.device] = None,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
) -> torch.Tensor:
|
||||
scale = (self.target_area / (height * width)) ** 0.5
|
||||
t = torch.arange(num_frames * nccl_info.sp_size, device=device, dtype=dtype)
|
||||
h = self._centers(-height * scale / 2, height * scale / 2, height, device, dtype)
|
||||
w = self._centers(-width * scale / 2, width * scale / 2, width, device, dtype)
|
||||
|
||||
grid_t, grid_h, grid_w = torch.meshgrid(t, h, w, indexing="ij")
|
||||
|
||||
positions = torch.stack([grid_t, grid_h, grid_w], dim=-1).view(-1, 3)
|
||||
return positions
|
||||
|
||||
def _create_rope(self, freqs: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
|
||||
freqs = torch.einsum("nd,dhf->nhf", pos, freqs.float())
|
||||
freqs_cos = torch.cos(freqs)
|
||||
freqs_sin = torch.sin(freqs)
|
||||
return freqs_cos, freqs_sin
|
||||
|
||||
def forward(
|
||||
self,
|
||||
pos_frequencies: torch.Tensor,
|
||||
num_frames: int,
|
||||
height: int,
|
||||
width: int,
|
||||
device: Optional[torch.device] = None,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
pos = self._get_positions(num_frames, height, width, device, dtype)
|
||||
rope_cos, rope_sin = self._create_rope(pos_frequencies, pos)
|
||||
return rope_cos, rope_sin
|
||||
|
||||
|
||||
@maybe_allow_in_graph
|
||||
class MochiTransformer3DModel(ModelMixin, ConfigMixin):
|
||||
r"""
|
||||
A Transformer model for video-like data introduced in [Mochi](https://huggingface.co/genmo/mochi-1-preview).
|
||||
|
||||
Args:
|
||||
patch_size (`int`, defaults to `2`):
|
||||
The size of the patches to use in the patch embedding layer.
|
||||
num_attention_heads (`int`, defaults to `24`):
|
||||
The number of heads to use for multi-head attention.
|
||||
attention_head_dim (`int`, defaults to `128`):
|
||||
The number of channels in each head.
|
||||
num_layers (`int`, defaults to `48`):
|
||||
The number of layers of Transformer blocks to use.
|
||||
in_channels (`int`, defaults to `12`):
|
||||
The number of channels in the input.
|
||||
out_channels (`int`, *optional*, defaults to `None`):
|
||||
The number of channels in the output.
|
||||
qk_norm (`str`, defaults to `"rms_norm"`):
|
||||
The normalization layer to use.
|
||||
text_embed_dim (`int`, defaults to `4096`):
|
||||
Input dimension of text embeddings from the text encoder.
|
||||
time_embed_dim (`int`, defaults to `256`):
|
||||
Output dimension of timestep embeddings.
|
||||
activation_fn (`str`, defaults to `"swiglu"`):
|
||||
Activation function to use in feed-forward.
|
||||
max_sequence_length (`int`, defaults to `256`):
|
||||
The maximum sequence length of text embeddings supported.
|
||||
"""
|
||||
|
||||
_supports_gradient_checkpointing = True
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
patch_size: int = 2,
|
||||
num_attention_heads: int = 24,
|
||||
attention_head_dim: int = 128,
|
||||
num_layers: int = 48,
|
||||
pooled_projection_dim: int = 1536,
|
||||
in_channels: int = 12,
|
||||
out_channels: Optional[int] = None,
|
||||
qk_norm: str = "rms_norm",
|
||||
text_embed_dim: int = 4096,
|
||||
time_embed_dim: int = 256,
|
||||
activation_fn: str = "swiglu",
|
||||
max_sequence_length: int = 256,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
inner_dim = num_attention_heads * attention_head_dim
|
||||
out_channels = out_channels or in_channels
|
||||
|
||||
self.patch_embed = PatchEmbed(
|
||||
patch_size=patch_size,
|
||||
in_channels=in_channels,
|
||||
embed_dim=inner_dim,
|
||||
pos_embed_type=None,
|
||||
)
|
||||
|
||||
self.time_embed = MochiCombinedTimestepCaptionEmbedding(
|
||||
embedding_dim=inner_dim,
|
||||
pooled_projection_dim=pooled_projection_dim,
|
||||
text_embed_dim=text_embed_dim,
|
||||
time_embed_dim=time_embed_dim,
|
||||
num_attention_heads=8,
|
||||
)
|
||||
|
||||
self.pos_frequencies = nn.Parameter(torch.full((3, num_attention_heads, attention_head_dim // 2), 0.0))
|
||||
self.rope = MochiRoPE()
|
||||
|
||||
self.transformer_blocks = nn.ModuleList(
|
||||
[
|
||||
MochiTransformerBlock(
|
||||
dim=inner_dim,
|
||||
num_attention_heads=num_attention_heads,
|
||||
attention_head_dim=attention_head_dim,
|
||||
pooled_projection_dim=pooled_projection_dim,
|
||||
qk_norm=qk_norm,
|
||||
activation_fn=activation_fn,
|
||||
context_pre_only=i == num_layers - 1,
|
||||
)
|
||||
for i in range(num_layers)
|
||||
]
|
||||
)
|
||||
|
||||
self.norm_out = AdaLayerNormContinuous(
|
||||
inner_dim, inner_dim, elementwise_affine=False, eps=1e-6, norm_type="layer_norm"
|
||||
)
|
||||
self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * out_channels)
|
||||
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def _set_gradient_checkpointing(self, module, value=False):
|
||||
if hasattr(module, "gradient_checkpointing"):
|
||||
module.gradient_checkpointing = value
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
timestep: torch.LongTensor,
|
||||
encoder_attention_mask: torch.Tensor,
|
||||
output_attn = False,
|
||||
return_dict: bool = False,
|
||||
) -> torch.Tensor:
|
||||
assert return_dict is False, "return_dict is not supported in MochiTransformer3DModel"
|
||||
batch_size, num_channels, num_frames, height, width = hidden_states.shape
|
||||
p = self.config.patch_size
|
||||
|
||||
post_patch_height = height // p
|
||||
post_patch_width = width // p
|
||||
|
||||
temb, encoder_hidden_states = self.time_embed(
|
||||
timestep, encoder_hidden_states, encoder_attention_mask, hidden_dtype=hidden_states.dtype
|
||||
)
|
||||
|
||||
hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1)
|
||||
hidden_states = self.patch_embed(hidden_states)
|
||||
hidden_states = hidden_states.unflatten(0, (batch_size, -1)).flatten(1, 2)
|
||||
|
||||
image_rotary_emb = self.rope(
|
||||
self.pos_frequencies,
|
||||
num_frames,
|
||||
post_patch_height,
|
||||
post_patch_width,
|
||||
device=hidden_states.device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
attn_outputs_list = []
|
||||
for i, block in enumerate(self.transformer_blocks):
|
||||
if self.gradient_checkpointing:
|
||||
|
||||
def create_custom_forward(module):
|
||||
def custom_forward(*inputs):
|
||||
return module(*inputs)
|
||||
|
||||
return custom_forward
|
||||
|
||||
ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
|
||||
hidden_states, encoder_hidden_states, attn_outputs = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(block),
|
||||
hidden_states,
|
||||
encoder_hidden_states,
|
||||
temb,
|
||||
image_rotary_emb,
|
||||
output_attn,
|
||||
**ckpt_kwargs,
|
||||
)
|
||||
else:
|
||||
hidden_states, encoder_hidden_states, attn_outputs = block(
|
||||
hidden_states=hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
temb=temb,
|
||||
image_rotary_emb=image_rotary_emb,
|
||||
output_attn = output_attn,
|
||||
)
|
||||
attn_outputs_list.append(attn_outputs)
|
||||
|
||||
hidden_states = self.norm_out(hidden_states, temb)
|
||||
hidden_states = self.proj_out(hidden_states)
|
||||
|
||||
hidden_states = hidden_states.reshape(batch_size, num_frames, post_patch_height, post_patch_width, p, p, -1)
|
||||
hidden_states = hidden_states.permute(0, 6, 1, 2, 4, 3, 5)
|
||||
output = hidden_states.reshape(batch_size, -1, num_frames, height, width)
|
||||
|
||||
if not output_attn :
|
||||
attn_outputs_list = None
|
||||
else:
|
||||
attn_outputs_list = torch.stack(attn_outputs_list, dim=0)
|
||||
return (output, attn_outputs_list)
|
||||
@@ -21,7 +21,8 @@ from transformers import T5EncoderModel, T5TokenizerFast
|
||||
|
||||
from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
|
||||
from diffusers.models.autoencoders import AutoencoderKL
|
||||
from diffusers.models.transformers import MochiTransformer3DModel
|
||||
from fastvideo.model.modeling_mochi import MochiTransformer3DModel
|
||||
|
||||
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
|
||||
from diffusers.utils import (
|
||||
is_torch_xla_available,
|
||||
|
||||
@@ -7,8 +7,8 @@ from fastvideo.utils.parallel_states import initialize_sequence_parallel_state,
|
||||
import argparse
|
||||
import os
|
||||
from diffusers.models.transformers.transformer_mochi import MochiTransformerBlock
|
||||
from fastvideo.model.mochi_monkey_patches import hf_mochi_add_sp_monkey_patch
|
||||
from diffusers import MochiTransformer3DModel
|
||||
from fastvideo.model.modeling_mochi import MochiTransformer3DModel
|
||||
|
||||
|
||||
|
||||
import sys
|
||||
@@ -123,7 +123,6 @@ if __name__ == "__main__":
|
||||
single_no_patch_bs_2 = test_single_block(2)
|
||||
# check all close
|
||||
assert torch.allclose(single_no_patch_bs_1, single_no_patch_bs_2)
|
||||
hf_mochi_add_sp_monkey_patch()
|
||||
single_patch_bs_1 = test_single_block(1)
|
||||
single_patch_bs_2 = test_single_block(2)
|
||||
assert torch.allclose(single_patch_bs_1, single_patch_bs_2)
|
||||
@@ -141,7 +140,6 @@ if __name__ == "__main__":
|
||||
single_no_patch_bs_1 = test_DiT(1, transformer, seed)
|
||||
single_no_patch_bs_2_a, single_no_patch_bs_2_b = test_DiT(2, transformer, seed)
|
||||
|
||||
hf_mochi_add_sp_monkey_patch()
|
||||
|
||||
single_patch_bs_1 = test_DiT(1, transformer, seed)
|
||||
single_patch_bs_2_a, single_patch_bs_2_b = test_DiT(2, transformer, seed)
|
||||
|
||||
@@ -6,8 +6,7 @@ from diffusers.utils import export_to_video
|
||||
from fastvideo.utils.parallel_states import initialize_sequence_parallel_state, nccl_info
|
||||
import argparse
|
||||
import os
|
||||
from diffusers import MochiTransformer3DModel
|
||||
from fastvideo.model.mochi_monkey_patches import hf_mochi_add_sp_monkey_patch
|
||||
from fastvideo.model.modeling_mochi import MochiTransformer3DModel
|
||||
|
||||
def initialize_distributed():
|
||||
local_rank = int(os.getenv('RANK', 0))
|
||||
@@ -21,7 +20,6 @@ def initialize_distributed():
|
||||
|
||||
def main(args):
|
||||
initialize_distributed()
|
||||
hf_mochi_add_sp_monkey_patch()
|
||||
print(nccl_info.sp_size)
|
||||
device = torch.cuda.current_device()
|
||||
generator = torch.Generator(device).manual_seed(args.seed)
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import torch
|
||||
from diffusers import MochiPipeline, MochiTransformer3DModel
|
||||
from fastvideo.model.pipeline_mochi import MochiPipeline
|
||||
from fastvideo.model.modeling_mochi import MochiTransformer3DModel
|
||||
from diffusers.utils import export_to_video, load_image, load_video
|
||||
import argparse
|
||||
|
||||
|
||||
+175
-401
@@ -4,39 +4,38 @@ import math
|
||||
import os
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from tqdm import tqdm
|
||||
from diffusers.training_utils import cast_training_params, compute_density_for_timestep_sampling, compute_loss_weighting_for_sd3
|
||||
from fastvideo.utils.parallel_states import initialize_sequence_parallel_state, \
|
||||
destroy_sequence_parallel_group, get_sequence_parallel_state, nccl_info
|
||||
from fastvideo.utils.communications import prepare_sequence_parallel_data, broadcast
|
||||
from fastvideo.model.mochi_monkey_patches import hf_mochi_add_sp_monkey_patch
|
||||
from fastvideo.utils.communications import sp_parallel_dataloader_wrapper, broadcast
|
||||
from fastvideo.model.mochi_latents_stat import mochi_stat
|
||||
from fastvideo.utils.validation import log_validation
|
||||
import time
|
||||
from torch.utils.data import DataLoader
|
||||
from copy import deepcopy
|
||||
import torch
|
||||
import transformers
|
||||
from accelerate import Accelerator, init_empty_weights, DistributedType
|
||||
from diffusers.utils.torch_utils import is_compiled_module
|
||||
from accelerate.logging import get_logger
|
||||
from accelerate.utils import DistributedType, ProjectConfiguration, set_seed
|
||||
from torch.distributed.fsdp import (
|
||||
FullyShardedDataParallel as FSDP,
|
||||
StateDictType,
|
||||
FullStateDictConfig,
|
||||
)
|
||||
import json
|
||||
from torch.utils.data.distributed import DistributedSampler
|
||||
import wandb
|
||||
from accelerate.utils import set_seed
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from fastvideo.fsdp_util import get_fsdp_kwargs, apply_fsdp_checkpointing
|
||||
import diffusers
|
||||
from diffusers import (
|
||||
FlowMatchEulerDiscreteScheduler,
|
||||
MochiTransformer3DModel,
|
||||
)
|
||||
from diffusers.optimization import get_scheduler
|
||||
from diffusers.utils import check_min_version, is_wandb_available
|
||||
from fastvideo.model.modeling_mochi import MochiTransformer3DModel
|
||||
from diffusers.utils import check_min_version
|
||||
from fastvideo.utils.ema import EMAModel
|
||||
from fastvideo.dataset.latent_datasets import LatentDataset, latent_collate_function
|
||||
|
||||
|
||||
import torch.distributed as dist
|
||||
from safetensors.torch import save_file
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.31.0")
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
import sys
|
||||
@@ -57,191 +56,178 @@ class ForkedPdb(pdb.Pdb):
|
||||
|
||||
|
||||
|
||||
def main_print(content):
|
||||
if int(os.environ['LOCAL_RANK']) <= 0:
|
||||
print(content)
|
||||
|
||||
def save_checkpoint(transformer: MochiTransformer3DModel, rank, output_dir, step):
|
||||
main_print(f"--> saving checkpoint at step {step}")
|
||||
with FSDP.state_dict_type(
|
||||
transformer, StateDictType.FULL_STATE_DICT, FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
|
||||
):
|
||||
cpu_state = transformer.state_dict()
|
||||
|
||||
if rank <= 0:
|
||||
save_dir = os.path.join(output_dir, f"checkpoint-{step}")
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
# save using safetensors
|
||||
weight_path = os.path.join(save_dir, "diffusion_pytorch_model.safetensors")
|
||||
save_file(cpu_state, weight_path)
|
||||
config_dict = dict(transformer.config)
|
||||
config_path = os.path.join(save_dir, "config.json")
|
||||
# save dict as json
|
||||
with open(config_path, "w", indent=4) as f:
|
||||
json.dump(config_dict, f)
|
||||
main_print(f"--> checkpoint saved at step {step}")
|
||||
|
||||
|
||||
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 ProgressInfo:
|
||||
def __init__(self, global_step, train_loss=0.0):
|
||||
self.global_step = global_step
|
||||
self.train_loss = train_loss
|
||||
def train_one_step_mochi(transformer, optimizer, loader,noise_scheduler, gradient_accumulation_steps, sp_size, precondition_outputs, max_grad_norm):
|
||||
total_loss = 0.0
|
||||
optimizer.zero_grad()
|
||||
for _ in range(gradient_accumulation_steps):
|
||||
latents, encoder_hidden_states, latents_attention_mask, encoder_attention_mask = next(loader)
|
||||
|
||||
latents = latents.to(dtype=torch.bfloat16)
|
||||
encoder_hidden_states = encoder_hidden_states.to(dtype=torch.bfloat16)
|
||||
latents_mean = (
|
||||
torch.tensor(mochi_stat.latents_mean).view(1, 12, 1, 1, 1).to(latents.device, latents.dtype)
|
||||
)
|
||||
latents_std = (
|
||||
torch.tensor(mochi_stat.latents_std).view(1, 12, 1, 1, 1).to(latents.device, latents.dtype)
|
||||
)
|
||||
latents = (latents - latents_mean) / latents_std
|
||||
|
||||
batch_size = latents.shape[0]
|
||||
noise = torch.randn_like(latents)
|
||||
u = torch.rand(size=(batch_size,), device="cpu")
|
||||
indices = (u * noise_scheduler.config.num_train_timesteps).long()
|
||||
timesteps = noise_scheduler.timesteps[indices].to(device=latents.device)
|
||||
if sp_size > 1:
|
||||
# Make sure that the timesteps are the same across all sp processes.
|
||||
broadcast(timesteps)
|
||||
|
||||
sigmas = get_sigmas(noise_scheduler, latents.device, timesteps, n_dim=latents.ndim, dtype=latents.dtype)
|
||||
noisy_model_input = (1.0 - sigmas) * latents + sigmas * noise
|
||||
|
||||
model_pred = transformer(
|
||||
noisy_model_input,
|
||||
encoder_hidden_states,
|
||||
noise_scheduler.config.num_train_timesteps - timesteps,
|
||||
encoder_attention_mask, # B, L
|
||||
return_dict= False
|
||||
)[0]
|
||||
|
||||
if precondition_outputs:
|
||||
model_pred = model_pred * sigmas + noisy_model_input
|
||||
|
||||
|
||||
#################################################################################
|
||||
# Training Loop #
|
||||
#################################################################################
|
||||
|
||||
|
||||
if precondition_outputs:
|
||||
target = latents
|
||||
else:
|
||||
target = latents - noise
|
||||
|
||||
loss = torch.mean((model_pred.float() - target.float()) ** 2)
|
||||
|
||||
loss.backward()
|
||||
|
||||
avg_loss = loss.detach().clone()
|
||||
dist.all_reduce(avg_loss, op=dist.ReduceOp.AVG)
|
||||
total_loss += avg_loss.item() / gradient_accumulation_steps
|
||||
|
||||
|
||||
transformer.clip_grad_norm_(max_grad_norm)
|
||||
optimizer.step()
|
||||
return total_loss
|
||||
|
||||
|
||||
def main(args):
|
||||
logging_dir = Path(args.output_dir, args.logging_dir)
|
||||
|
||||
# use LayerNorm, GeLu, SiLu always as fp32 mode
|
||||
# TODO:
|
||||
if args.enable_stable_fp32:
|
||||
raise NotImplementedError("enable_stable_fp32 is not supported now.")
|
||||
|
||||
accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
|
||||
|
||||
accelerator = Accelerator(
|
||||
gradient_accumulation_steps=args.gradient_accumulation_steps,
|
||||
mixed_precision=args.mixed_precision,
|
||||
log_with="wandb",
|
||||
project_config=accelerator_project_config,
|
||||
)
|
||||
|
||||
initialize_sequence_parallel_state(args.sp_size)
|
||||
hf_mochi_add_sp_monkey_patch()
|
||||
torch.backends.cuda.matmul.allow_tf32 = True
|
||||
|
||||
if not is_wandb_available():
|
||||
raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
|
||||
# Make one log on every process with the configuration for debugging.
|
||||
logging.basicConfig(
|
||||
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
|
||||
datefmt="%m/%d/%Y %H:%M:%S",
|
||||
level=logging.INFO,
|
||||
)
|
||||
logger.info(accelerator.state, main_process_only=False)
|
||||
if accelerator.is_local_main_process:
|
||||
transformers.utils.logging.set_verbosity_warning()
|
||||
diffusers.utils.logging.set_verbosity_info()
|
||||
else:
|
||||
transformers.utils.logging.set_verbosity_error()
|
||||
diffusers.utils.logging.set_verbosity_error()
|
||||
local_rank = int(os.environ['LOCAL_RANK'])
|
||||
rank = int(os.environ['RANK'])
|
||||
world_size = int(os.environ['WORLD_SIZE'])
|
||||
dist.init_process_group("nccl")
|
||||
torch.cuda.set_device(local_rank)
|
||||
device = torch.cuda.current_device()
|
||||
initialize_sequence_parallel_state(args.sp_size)
|
||||
|
||||
|
||||
|
||||
|
||||
# If passed along, set the training seed now. On GPU...
|
||||
if args.seed is not None:
|
||||
# TODO: t within the same seq parallel group should be the same. Noise should be different.
|
||||
set_seed(args.seed + accelerator.process_index)
|
||||
set_seed(args.seed + rank)
|
||||
|
||||
# Handle the repository creation
|
||||
if accelerator.is_main_process:
|
||||
if args.output_dir is not None:
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
if rank <=0 and args.output_dir is not None:
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
# For mixed precision training we cast all non-trainable weigths to half-precision
|
||||
# as these weights are only used for inference, keeping weights in full precision is not required.
|
||||
weight_dtype = torch.float32
|
||||
if accelerator.mixed_precision == "fp16":
|
||||
weight_dtype = torch.float16
|
||||
elif accelerator.mixed_precision == "bf16":
|
||||
weight_dtype = torch.bfloat16
|
||||
|
||||
|
||||
# Create model:
|
||||
|
||||
|
||||
load_dtype = torch.bfloat16 # TODO
|
||||
main_print(f"--> loading model from {args.pretrained_model_name_or_path}")
|
||||
transformer = MochiTransformer3DModel.from_pretrained(
|
||||
args.pretrained_model_name_or_path,
|
||||
subfolder="transformer",
|
||||
torch_dtype=load_dtype,
|
||||
)
|
||||
|
||||
|
||||
main_print(f"--> Initializing FSDP with sharding strategy: full")
|
||||
fsdp_kwargs = get_fsdp_kwargs("full", args.use_cpu_offload)
|
||||
transformer = FSDP(
|
||||
transformer,
|
||||
**fsdp_kwargs,
|
||||
)
|
||||
|
||||
main_print(f"--> model loaded")
|
||||
if args.gradient_checkpointing:
|
||||
transformer.enable_gradient_checkpointing()
|
||||
|
||||
apply_fsdp_checkpointing(transformer)
|
||||
|
||||
|
||||
# Set model as trainable.
|
||||
transformer.train()
|
||||
|
||||
noise_scheduler = FlowMatchEulerDiscreteScheduler()
|
||||
|
||||
|
||||
# Create EMA for the unet.
|
||||
if args.use_ema:
|
||||
ema_transformer = deepcopy(transformer)
|
||||
ema_transformer = EMAModel(ema_transformer.parameters(), decay=args.ema_decay, update_after_step=args.ema_start_step,
|
||||
model_cls=MochiTransformer3DModel, model_config=ema_transformer.config)
|
||||
|
||||
def unwrap_model(model):
|
||||
model = accelerator.unwrap_model(model)
|
||||
model = model._orig_mod if is_compiled_module(model) else model
|
||||
return model
|
||||
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
|
||||
def save_model_hook(models, weights, output_dir):
|
||||
if accelerator.is_main_process:
|
||||
if args.use_ema:
|
||||
ema_transformer.save_pretrained(os.path.join(output_dir, "model_ema"))
|
||||
|
||||
for model in models:
|
||||
if isinstance(unwrap_model(model), type(unwrap_model(transformer))):
|
||||
model: MochiTransformer3DModel = unwrap_model(model)
|
||||
model.save_pretrained(
|
||||
os.path.join(output_dir, "transformer"), safe_serialization=True, max_shard_size="5GB"
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unexpected save model: {model.__class__}")
|
||||
# make sure to pop weight so that corresponding model is not saved again
|
||||
if weights:
|
||||
weights.pop()
|
||||
|
||||
|
||||
def load_model_hook(models, input_dir):
|
||||
if args.use_ema:
|
||||
load_model = EMAModel.from_pretrained(os.path.join(input_dir, "model_ema"), MochiTransformer3DModel)
|
||||
ema_transformer.load_state_dict(load_model.state_dict())
|
||||
ema_transformer.to(accelerator.device)
|
||||
del load_model
|
||||
|
||||
if not accelerator.distributed_type == DistributedType.DEEPSPEED:
|
||||
while len(models) > 0:
|
||||
model = models.pop()
|
||||
|
||||
if isinstance(unwrap_model(model), type(unwrap_model(transformer))):
|
||||
transformer_ = unwrap_model(model)
|
||||
else:
|
||||
raise ValueError(f"Unexpected save model: {unwrap_model(model).__class__}")
|
||||
else:
|
||||
with init_empty_weights():
|
||||
transformer_ = MochiTransformer3DModel.from_config(
|
||||
args.pretrained_model_name_or_path, subfolder="transformer"
|
||||
)
|
||||
init_under_meta = True
|
||||
|
||||
load_model = MochiTransformer3DModel.from_pretrained(os.path.join(input_dir, "transformer"))
|
||||
transformer_.register_to_config(**load_model.config)
|
||||
transformer_.load_state_dict(load_model.state_dict(), assign=init_under_meta)
|
||||
del load_model
|
||||
|
||||
if args.mixed_precision == "fp16":
|
||||
cast_training_params([transformer_])
|
||||
|
||||
accelerator.register_save_state_pre_hook(save_model_hook)
|
||||
accelerator.register_load_state_pre_hook(load_model_hook)
|
||||
|
||||
# Enable TF32 for faster training on Ampere GPUs,
|
||||
# cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
|
||||
if args.allow_tf32 and torch.cuda.is_available():
|
||||
torch.backends.cuda.matmul.allow_tf32 = True
|
||||
|
||||
if args.scale_lr:
|
||||
args.learning_rate = (
|
||||
args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes
|
||||
)
|
||||
|
||||
# Make sure the trainable params are in float32.
|
||||
# if args.mixed_precision == "fp16" or args.mixed_precision == "bf16":
|
||||
# cast_training_params([transformer], dtype=torch.float32)
|
||||
|
||||
|
||||
params_to_optimize = transformer.parameters()
|
||||
|
||||
optimizer = torch.optim.AdamW(
|
||||
params_to_optimize,
|
||||
transformer.parameters(),
|
||||
lr=args.learning_rate,
|
||||
betas=(0.9,0.999),
|
||||
weight_decay=0.01,
|
||||
eps=1e-8,
|
||||
)
|
||||
|
||||
lr_scheduler = get_scheduler(
|
||||
"constant",
|
||||
optimizer=optimizer,
|
||||
num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
|
||||
num_training_steps=args.max_train_steps * accelerator.num_processes,
|
||||
)
|
||||
|
||||
logger.info(f"optimizer: {optimizer}")
|
||||
main_print(f"optimizer: {optimizer}")
|
||||
|
||||
train_dataset = LatentDataset(args.data_json_path, args.num_latent_t, args.cfg)
|
||||
sampler = DistributedSampler(train_dataset, rank=rank, num_replicas=world_size, shuffle=True)
|
||||
train_dataloader = DataLoader(
|
||||
train_dataset,
|
||||
shuffle=True,
|
||||
sampler=sampler,
|
||||
collate_fn=latent_collate_function,
|
||||
pin_memory=True,
|
||||
batch_size=args.train_batch_size,
|
||||
@@ -251,17 +237,6 @@ def main(args):
|
||||
|
||||
|
||||
|
||||
# Prepare everything with our `accelerator`.
|
||||
# model.requires_grad_(False)
|
||||
# model.pos_embed.requires_grad_(True)
|
||||
logger.info(f'before accelerator.prepare')
|
||||
transformer, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
|
||||
transformer, optimizer, train_dataloader, lr_scheduler
|
||||
)
|
||||
logger.info(f'after accelerator.prepare')
|
||||
if args.use_ema:
|
||||
ema_transformer.to(accelerator.device)
|
||||
|
||||
# We need to recalculate our total training steps as the size of the training dataloader may have changed.
|
||||
num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps * args.sp_size / args.train_sp_batch_size)
|
||||
# Afterwards we recalculate our number of training epochs
|
||||
@@ -269,237 +244,56 @@ def main(args):
|
||||
|
||||
# We need to initialize the trackers we use, and also store our configuration.
|
||||
# The trackers initializes automatically on the main process.
|
||||
if accelerator.is_main_process:
|
||||
tracker_project_name = args.tracker_project_name or "fastvideo"
|
||||
accelerator.init_trackers(tracker_project_name, config=vars(args))
|
||||
if rank <= 0:
|
||||
project = args.tracker_project_name or "fastvideo"
|
||||
wandb.init(project=project, config=args)
|
||||
|
||||
# Train!
|
||||
total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps / args.sp_size * args.train_sp_batch_size
|
||||
logger.info("***** Running training *****")
|
||||
logger.info(f" Num examples = {len(train_dataset)}")
|
||||
logger.info(f" Dataloader size = {len(train_dataloader)}")
|
||||
logger.info(f" Num Epochs = {args.num_train_epochs}")
|
||||
logger.info(f" Instantaneous batch size per device = {args.train_batch_size}")
|
||||
logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
|
||||
logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}")
|
||||
logger.info(f" Total optimization steps = {args.max_train_steps}")
|
||||
logger.info(f" Total training parameters = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e9} B")
|
||||
global_step = 0
|
||||
first_epoch = 0
|
||||
total_batch_size = args.train_batch_size * world_size * args.gradient_accumulation_steps / args.sp_size * args.train_sp_batch_size
|
||||
main_print("***** Running training *****")
|
||||
main_print(f" Num examples = {len(train_dataset)}")
|
||||
main_print(f" Dataloader size = {len(train_dataloader)}")
|
||||
main_print(f" Num Epochs = {args.num_train_epochs}")
|
||||
main_print(f" Instantaneous batch size per device = {args.train_batch_size}")
|
||||
main_print(f" Total train batch size (w. data & sequence parallel, accumulation) = {total_batch_size}")
|
||||
main_print(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}")
|
||||
main_print(f" Total optimization steps = {args.max_train_steps}")
|
||||
main_print(f" Total training parameters per FSDP shard = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e9} B")
|
||||
# print dtype
|
||||
main_print(f" Master weight dtype: {transformer.parameters().__next__().dtype}")
|
||||
|
||||
|
||||
# Potentially load in the weights and states from a previous save
|
||||
if args.resume_from_checkpoint:
|
||||
if args.resume_from_checkpoint != "latest":
|
||||
path = os.path.basename(args.resume_from_checkpoint)
|
||||
else:
|
||||
# Get the most recent checkpoint
|
||||
dirs = os.listdir(args.output_dir)
|
||||
dirs = [d for d in dirs if d.startswith("checkpoint")]
|
||||
dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
|
||||
path = dirs[-1] if len(dirs) > 0 else None
|
||||
assert NotImplementedError("resume_from_checkpoint is not supported now.")
|
||||
# TODO
|
||||
|
||||
if path is None:
|
||||
accelerator.print(
|
||||
f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
|
||||
)
|
||||
args.resume_from_checkpoint = None
|
||||
initial_global_step = 0
|
||||
else:
|
||||
accelerator.print(f"Resuming from checkpoint {path}")
|
||||
accelerator.load_state(os.path.join(args.output_dir, path))
|
||||
global_step = int(path.split("-")[1])
|
||||
|
||||
initial_global_step = global_step
|
||||
first_epoch = global_step // num_update_steps_per_epoch
|
||||
else:
|
||||
initial_global_step = 0
|
||||
|
||||
progress_bar = tqdm(
|
||||
range(0, args.max_train_steps),
|
||||
initial=initial_global_step,
|
||||
initial=0,
|
||||
desc="Steps",
|
||||
# Only show the progress bar once on each machine.
|
||||
disable=not accelerator.is_local_main_process,
|
||||
disable= local_rank > 0,
|
||||
)
|
||||
progress_info = ProgressInfo(global_step, train_loss=0.0)
|
||||
|
||||
def sync_gradients_info():
|
||||
# Checks if the accelerator has performed an optimization step behind the scenesv
|
||||
if args.use_ema:
|
||||
ema_transformer.step(transformer.parameters())
|
||||
progress_bar.update(1)
|
||||
progress_info.global_step += 1
|
||||
end_time = time.time()
|
||||
one_step_duration = end_time - start_time
|
||||
accelerator.log(
|
||||
{
|
||||
"train_loss": progress_info.train_loss,
|
||||
"learning_rate": lr_scheduler.get_last_lr()[0]
|
||||
},
|
||||
step=progress_info.global_step
|
||||
)
|
||||
logs = {"step_loss": progress_info.train_loss, "lr": lr_scheduler.get_last_lr()[0]}
|
||||
progress_bar.set_postfix(**logs)
|
||||
|
||||
progress_info.train_loss = 0.0
|
||||
|
||||
# DeepSpeed requires saving weights on every device; saving weights only on the main process would cause issues.
|
||||
if accelerator.distributed_type == DistributedType.DEEPSPEED or accelerator.is_main_process:
|
||||
if progress_info.global_step % args.checkpointing_steps == 0:
|
||||
# _before_ saving state, check if this save would set us over the `checkpoints_total_limit`
|
||||
if accelerator.is_main_process and args.checkpoints_total_limit is not None:
|
||||
checkpoints = os.listdir(args.output_dir)
|
||||
checkpoints = [d for d in checkpoints if d.startswith("checkpoint")]
|
||||
checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1]))
|
||||
|
||||
# before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints
|
||||
if len(checkpoints) >= args.checkpoints_total_limit:
|
||||
num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1
|
||||
removing_checkpoints = checkpoints[0:num_to_remove]
|
||||
|
||||
logger.info(
|
||||
f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints"
|
||||
)
|
||||
logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}")
|
||||
|
||||
for removing_checkpoint in removing_checkpoints:
|
||||
removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint)
|
||||
shutil.rmtree(removing_checkpoint)
|
||||
|
||||
save_path = os.path.join(args.output_dir, f"checkpoint-{progress_info.global_step}")
|
||||
accelerator.save_state(save_path)
|
||||
logger.info(f"Saved state to {save_path}")
|
||||
|
||||
def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
|
||||
sigmas = noise_scheduler.sigmas.to(device=accelerator.device, dtype=dtype)
|
||||
schedule_timesteps = noise_scheduler.timesteps.to(accelerator.device)
|
||||
timesteps = timesteps.to(accelerator.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 run(latents, encoder_hidden_states,encoder_attention_mask):
|
||||
global start_time
|
||||
start_time = time.time()
|
||||
bsz = latents.shape[0]
|
||||
noise = torch.randn_like(latents)
|
||||
u = compute_density_for_timestep_sampling(
|
||||
weighting_scheme=args.weighting_scheme,
|
||||
batch_size=bsz,
|
||||
logit_mean=args.logit_mean,
|
||||
logit_std=args.logit_std,
|
||||
mode_scale=args.mode_scale,
|
||||
)
|
||||
indices = (u * noise_scheduler.config.num_train_timesteps).long()
|
||||
timesteps = noise_scheduler.timesteps[indices].to(device=latents.device)
|
||||
if args.sp_size > 1:
|
||||
# Make sure that the timesteps are the same across all sp processes.
|
||||
broadcast(timesteps)
|
||||
# Add noise according to flow matching.
|
||||
# zt = (1 - texp) * x + texp * z1
|
||||
sigmas = get_sigmas(timesteps, n_dim=latents.ndim, dtype=latents.dtype)
|
||||
noisy_model_input = (1.0 - sigmas) * latents + sigmas * noise
|
||||
# accelerator.print(latents.shape, encoder_hidden_states.shape, encoder_attention_mask.shape)
|
||||
# accelerator.print(encoder_attention_mask.tolist())
|
||||
# print(timesteps)
|
||||
# accelerator.print(sigmas, timesteps)
|
||||
model_pred = transformer(
|
||||
noisy_model_input,
|
||||
encoder_hidden_states,
|
||||
noise_scheduler.config.num_train_timesteps - timesteps,
|
||||
encoder_attention_mask, # B, L
|
||||
return_dict= False
|
||||
)[0]
|
||||
|
||||
if args.precondition_outputs:
|
||||
model_pred = model_pred * sigmas + noisy_model_input
|
||||
|
||||
# these weighting schemes use a uniform timestep sampling
|
||||
# and instead post-weight the loss
|
||||
weighting = compute_loss_weighting_for_sd3(weighting_scheme=args.weighting_scheme, sigmas=sigmas)
|
||||
|
||||
# flow matching loss
|
||||
if args.precondition_outputs:
|
||||
target = latents
|
||||
else:
|
||||
target = latents - noise
|
||||
|
||||
loss = torch.mean(weighting.float() * (model_pred.float() - target.float()) ** 2)
|
||||
|
||||
# Gather the losses across all processes for logging (if we use distributed training).
|
||||
# TODO Why repeat here? Weird
|
||||
avg_loss = accelerator.reduce(loss.clone().detach(), "mean")
|
||||
# accelerator.print(model_pred.shape)
|
||||
# accelerator.print(avg_loss)
|
||||
progress_info.train_loss += avg_loss.item() / args.gradient_accumulation_steps
|
||||
# Backpropagate
|
||||
accelerator.backward(loss)
|
||||
if accelerator.sync_gradients:
|
||||
accelerator.clip_grad_norm_(transformer.parameters(), args.max_grad_norm)
|
||||
optimizer.step()
|
||||
lr_scheduler.step()
|
||||
optimizer.zero_grad()
|
||||
|
||||
if accelerator.sync_gradients:
|
||||
# accelerator.print("Syncing gradients")
|
||||
sync_gradients_info()
|
||||
|
||||
# TODO: add indent
|
||||
#Group ID indicates which SP group the process belongs to. We only need SP group 0 to log validation.
|
||||
if progress_info.global_step % args.validation_steps == 0 and args.log_validation:
|
||||
log_validation(args, transformer, accelerator,
|
||||
weight_dtype, progress_info.global_step)
|
||||
if args.use_ema:
|
||||
# Store the UNet parameters temporarily and load the EMA parameters to perform inference.
|
||||
ema_transformer.store(transformer.parameters())
|
||||
ema_transformer.copy_to(transformer.parameters())
|
||||
log_validation(args, transformer, accelerator,
|
||||
weight_dtype, progress_info.global_step, ema=True)
|
||||
# Switch back to the original UNet parameters.
|
||||
ema_transformer.restore(transformer.parameters())
|
||||
|
||||
|
||||
loader = sp_parallel_dataloader_wrapper(train_dataloader, device, args.train_batch_size, args.sp_size, args.train_sp_batch_size)
|
||||
|
||||
for step in range(1, args.max_train_steps+1):
|
||||
loss = train_one_step_mochi(transformer, optimizer, loader, noise_scheduler, args.gradient_accumulation_steps, args.sp_size, args.precondition_outputs, args.max_grad_norm)
|
||||
progress_bar.set_postfix({"loss": loss})
|
||||
progress_bar.update(1)
|
||||
if rank <= 0:
|
||||
wandb.log({"train_loss": loss}, step=step)
|
||||
if step % args.checkpointing_steps == 0:
|
||||
save_checkpoint(transformer, rank, args.output_dir, step)
|
||||
|
||||
if args.log_validation and step % args.validation_steps == 0:
|
||||
log_validation(args, transformer, device,
|
||||
torch.bfloat16, step)
|
||||
|
||||
return loss
|
||||
|
||||
def sp_parallel_dataloader_wrapper(dataloader):
|
||||
for data_item in dataloader:
|
||||
latents, cond,attn_mask, cond_mask = data_item
|
||||
latents, cond, attn_mask, cond_mask = prepare_sequence_parallel_data(latents, cond, attn_mask, cond_mask)
|
||||
# accelerator.print(latents.shape, cond.shape, attn_mask.shape, cond_mask.shape)
|
||||
for iter in range(args.train_batch_size * args.sp_size // args.train_sp_batch_size):
|
||||
st_idx = iter * args.train_sp_batch_size
|
||||
ed_idx = (iter + 1) * args.train_sp_batch_size
|
||||
encoder_hidden_states=cond[st_idx: ed_idx]
|
||||
attention_mask=attn_mask[st_idx: ed_idx]
|
||||
encoder_attention_mask=cond_mask[st_idx: ed_idx]
|
||||
yield latents[st_idx: ed_idx], encoder_hidden_states, attention_mask, encoder_attention_mask
|
||||
|
||||
|
||||
def train():
|
||||
for epoch in range(first_epoch, args.num_train_epochs):
|
||||
progress_info.train_loss = 0.0
|
||||
for latents, cond, attn_mask, cond_mask in sp_parallel_dataloader_wrapper(train_dataloader):
|
||||
latents = latents.to(dtype=torch.bfloat16)
|
||||
cond = cond.to(dtype=torch.bfloat16)
|
||||
latents_mean = (
|
||||
torch.tensor(mochi_stat.latents_mean).view(1, 12, 1, 1, 1).to(latents.device, latents.dtype)
|
||||
)
|
||||
latents_std = (
|
||||
torch.tensor(mochi_stat.latents_std).view(1, 12, 1, 1, 1).to(latents.device, latents.dtype)
|
||||
)
|
||||
latents = (latents - latents_mean) / latents_std
|
||||
with accelerator.accumulate(transformer):
|
||||
run(latents, cond, cond_mask)
|
||||
|
||||
if progress_info.global_step >= args.max_train_steps:
|
||||
return
|
||||
|
||||
train()
|
||||
accelerator.wait_for_everyone()
|
||||
accelerator.end_training()
|
||||
if get_sequence_parallel_state():
|
||||
destroy_sequence_parallel_group()
|
||||
|
||||
@@ -522,30 +316,9 @@ if __name__ == "__main__":
|
||||
parser.add_argument('--enable_stable_fp32', action='store_true') # TODO
|
||||
|
||||
# diffusion setting
|
||||
parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.")
|
||||
parser.add_argument("--ema_decay", type=float, default=0.999)
|
||||
parser.add_argument("--ema_start_step", type=int, default=0)
|
||||
parser.add_argument('--cfg', type=float, default=0.1)
|
||||
parser.add_argument("--prediction_type", type=str, default=None, help="The prediction_type that shall be used for training. Choose between 'epsilon' or 'v_prediction' or leave `None`. If left to `None` the default prediction type of the scheduler: `noise_scheduler.config.prediciton_type` is chosen.")
|
||||
|
||||
parser.add_argument(
|
||||
"--weighting_scheme",
|
||||
type=str,
|
||||
default="uniform",
|
||||
choices=["sigma_sqrt", "logit_normal", "mode", "cosmap", "uniform"],
|
||||
)
|
||||
parser.add_argument(
|
||||
"--logit_mean", type=float, default=0.0, help="mean to use when using the `'logit_normal'` weighting scheme."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--logit_std", type=float, default=1.0, help="std to use when using the `'logit_normal'` weighting scheme."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mode_scale",
|
||||
type=float,
|
||||
default=1.29,
|
||||
help="Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.",
|
||||
)
|
||||
parser.add_argument("--precondition_outputs", action="store_true", help="Whether to precondition the outputs of the model.")
|
||||
|
||||
# validation & logs
|
||||
@@ -601,6 +374,7 @@ if __name__ == "__main__":
|
||||
" flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
|
||||
),
|
||||
)
|
||||
parser.add_argument("--use_cpu_offload", action="store_true", help="Whether to use CPU offload for param & gradient & optimizer states.")
|
||||
|
||||
parser.add_argument("--sp_size", type=int, default=1, help="For sequence parallel")
|
||||
parser.add_argument("--train_sp_batch_size", type=int, default=1, help="Batch size for sequence parallel training")
|
||||
|
||||
@@ -274,4 +274,21 @@ def prepare_sequence_parallel_data(hidden_states, encoder_hidden_states, attenti
|
||||
return hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask
|
||||
|
||||
|
||||
|
||||
|
||||
def sp_parallel_dataloader_wrapper(dataloader, device, train_batch_size, sp_size, train_sp_batch_size):
|
||||
while True:
|
||||
for data_item in dataloader:
|
||||
latents, cond,attn_mask, cond_mask = data_item
|
||||
latents = latents.to(device)
|
||||
cond = cond.to(device)
|
||||
attn_mask = attn_mask.to(device)
|
||||
cond_mask = cond_mask.to(device)
|
||||
latents, cond, attn_mask, cond_mask = prepare_sequence_parallel_data(latents, cond, attn_mask, cond_mask)
|
||||
|
||||
for iter in range(train_batch_size * sp_size // train_sp_batch_size):
|
||||
st_idx = iter * train_sp_batch_size
|
||||
ed_idx = (iter + 1) * train_sp_batch_size
|
||||
encoder_hidden_states=cond[st_idx: ed_idx]
|
||||
attention_mask=attn_mask[st_idx: ed_idx]
|
||||
encoder_attention_mask=cond_mask[st_idx: ed_idx]
|
||||
yield latents[st_idx: ed_idx], encoder_hidden_states, attention_mask, encoder_attention_mask
|
||||
@@ -1,7 +1,7 @@
|
||||
import argparse
|
||||
import torch
|
||||
from accelerate.logging import get_logger
|
||||
from diffusers import MochiPipeline
|
||||
from fastvideo.model.pipeline_mochi import MochiPipeline
|
||||
from diffusers.utils import export_to_video
|
||||
import json
|
||||
import os
|
||||
|
||||
@@ -38,9 +38,9 @@ import diffusers
|
||||
from diffusers import (
|
||||
AutoencoderKL,
|
||||
FlowMatchEulerDiscreteScheduler,
|
||||
MochiTransformer3DModel,
|
||||
MochiPipeline,
|
||||
)
|
||||
from fastvideo.model.modeling_mochi import MochiTransformer3DModel
|
||||
from fastvideo.model.pipeline_mochi import MochiPipeline
|
||||
from diffusers.optimization import get_scheduler
|
||||
from diffusers.training_utils import compute_snr
|
||||
from diffusers.utils import check_min_version, is_wandb_available
|
||||
|
||||
@@ -177,12 +177,11 @@ def sample_validation_video(
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def log_validation(args, transformer,accelerator, weight_dtype, global_step, ema=False):
|
||||
def log_validation(args, transformer, device, weight_dtype, global_step, ema=False):
|
||||
#TODO
|
||||
accelerator.print(f"Running validation....\n")
|
||||
transformer = accelerator.unwrap_model(transformer)
|
||||
generator = torch.Generator(device=accelerator.device).manual_seed(12345)
|
||||
vae = AutoencoderKLMochi.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", torch_dtype=weight_dtype).to(accelerator.device)
|
||||
print(f"Running validation....\n")
|
||||
generator = torch.Generator(device="cuda").manual_seed(12345)
|
||||
vae = AutoencoderKLMochi.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", torch_dtype=weight_dtype).to("cuda")
|
||||
vae.enable_tiling()
|
||||
scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler", torch_dtype=weight_dtype)
|
||||
|
||||
@@ -191,10 +190,10 @@ def log_validation(args, transformer,accelerator, weight_dtype, global_step, em
|
||||
prompt_mask_path = os.path.join(args.validation_prompt_dir, "mask.pt")
|
||||
negative_prompt_embed_path = os.path.join(args.uncond_prompt_dir, "embed.pt")
|
||||
negative_prompt_mask_path = os.path.join(args.uncond_prompt_dir, "mask.pt")
|
||||
prompt_embeds = torch.load(prompt_embed_path, map_location=accelerator.device).to(accelerator.device).to(weight_dtype).unsqueeze(0)
|
||||
prompt_attention_mask = torch.load(prompt_mask_path, map_location=accelerator.device).to(accelerator.device).to(weight_dtype).unsqueeze(0)
|
||||
negative_prompt_embeds = torch.load(negative_prompt_embed_path, map_location=accelerator.device).to(accelerator.device).to(weight_dtype).unsqueeze(0)
|
||||
negative_prompt_attention_mask = torch.load(negative_prompt_mask_path, map_location=accelerator.device).to(accelerator.device).to(weight_dtype).unsqueeze(0)
|
||||
prompt_embeds = torch.load(prompt_embed_path, map_location="cpu", weights_only=True).to(device).to(weight_dtype).unsqueeze(0)
|
||||
prompt_attention_mask = torch.load(prompt_mask_path, map_location="cpu", weights_only=True).to(device).to(weight_dtype).unsqueeze(0)
|
||||
negative_prompt_embeds = torch.load(negative_prompt_embed_path, map_location="cpu", weights_only=True).to(device).to(weight_dtype).unsqueeze(0)
|
||||
negative_prompt_attention_mask = torch.load(negative_prompt_mask_path, map_location="cpu", weights_only=True).to(device).to(weight_dtype).unsqueeze(0)
|
||||
|
||||
|
||||
videos = []
|
||||
@@ -219,20 +218,20 @@ def log_validation(args, transformer,accelerator, weight_dtype, global_step, em
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
# log if main process
|
||||
if accelerator.is_main_process :
|
||||
if int(os.environ['RANK']) <= 0:
|
||||
video_filenames = []
|
||||
for i, video in enumerate(videos):
|
||||
filename = os.path.join(args.output_dir, f"validation_step_{global_step}_video_{i}.mp4")
|
||||
export_to_video(video, filename, fps=30)
|
||||
video_filenames.append(filename)
|
||||
for tracker in accelerator.trackers:
|
||||
logs = {
|
||||
f"{'ema_' if ema else ''}validation": [
|
||||
wandb.Video(filename)
|
||||
for i, filename in enumerate(video_filenames)
|
||||
]
|
||||
}
|
||||
tracker.log(logs, step=global_step)
|
||||
|
||||
logs = {
|
||||
f"{'ema_' if ema else ''}validation": [
|
||||
wandb.Video(filename)
|
||||
for i, filename in enumerate(video_filenames)
|
||||
]
|
||||
}
|
||||
wandb.log(logs, step=global_step)
|
||||
|
||||
del vae
|
||||
del prompt_embeds
|
||||
|
||||
+5
-5
@@ -17,18 +17,18 @@ torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
|
||||
|
||||
|
||||
|
||||
num_gpus=2
|
||||
num_gpus=4
|
||||
torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
|
||||
fastvideo/sample/sample_t2v_mochi.py \
|
||||
--model_path data/mochi \
|
||||
--prompt_embed_path data/Encoder_Overfit_Data/prompt_embed/1.pt \
|
||||
--encoder_attention_mask_path data/Encoder_Overfit_Data/prompt_attention_mask/1.pt \
|
||||
--num_frames 79 \
|
||||
--prompt_embed_path data/Encoder_Overfit_Data/prompt_embed/0.pt \
|
||||
--encoder_attention_mask_path data/Encoder_Overfit_Data/prompt_attention_mask/0.pt \
|
||||
--num_frames 163 \
|
||||
--height 480 \
|
||||
--width 848 \
|
||||
--num_inference_steps 64 \
|
||||
--guidance_scale 4.5 \
|
||||
--transformer_path data/outputs/debug_uniform_14_sp2_spb2_4videos_2K/checkpoint-1000/transformer \
|
||||
--transformer_path data/outputs/BW_Testrun/checkpoint-0 \
|
||||
--output_path outputs_video/overfit/uniform_cfg_train_1000_cfg_f79_sp2_spb2_4videos_second \
|
||||
--seed 12345
|
||||
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
import json
|
||||
import os
|
||||
|
||||
path = "data/Mochi-Synthetic-Data/videos2caption.json"
|
||||
path = "data/outputs/BW_Testrun/checkpoint-0/config.json"
|
||||
|
||||
with open(path, 'r') as f:
|
||||
data = json.load(f)
|
||||
|
||||
@@ -1,13 +0,0 @@
|
||||
compute_environment: LOCAL_MACHINE
|
||||
distributed_type: DEEPSPEED
|
||||
deepspeed_config:
|
||||
deepspeed_config_file: scripts/accelerate_configs/zero2.json
|
||||
fsdp_config: {}
|
||||
machine_rank: 0
|
||||
main_process_ip: null
|
||||
main_process_port: 29503
|
||||
main_training_function: main
|
||||
num_machines: 1
|
||||
num_processes: 4
|
||||
# gpu_ids: 0
|
||||
use_cpu: false
|
||||
@@ -1,13 +0,0 @@
|
||||
compute_environment: LOCAL_MACHINE
|
||||
distributed_type: DEEPSPEED
|
||||
deepspeed_config:
|
||||
deepspeed_config_file: scripts/accelerate_configs/zero2_offload.json
|
||||
fsdp_config: {}
|
||||
machine_rank: 0
|
||||
main_process_ip: null
|
||||
main_process_port: 29501
|
||||
main_training_function: main
|
||||
num_machines: 1
|
||||
num_processes: 8
|
||||
gpu_ids: 0,1,2,3,4,5,6,7
|
||||
use_cpu: false
|
||||
@@ -1,13 +0,0 @@
|
||||
compute_environment: LOCAL_MACHINE
|
||||
distributed_type: DEEPSPEED
|
||||
deepspeed_config:
|
||||
deepspeed_config_file: scripts/accelerate_configs/zero3.json
|
||||
zero3_init_flag: true
|
||||
fsdp_config: {}
|
||||
machine_rank: 0
|
||||
main_process_ip: null
|
||||
main_process_port: 29501
|
||||
main_training_function: main
|
||||
num_machines: 1
|
||||
num_processes: 4
|
||||
use_cpu: false
|
||||
@@ -1,13 +0,0 @@
|
||||
compute_environment: LOCAL_MACHINE
|
||||
distributed_type: DEEPSPEED
|
||||
deepspeed_config:
|
||||
deepspeed_config_file: scripts/accelerate_configs/zero3_offload.json
|
||||
fsdp_config: {}
|
||||
machine_rank: 0
|
||||
main_process_ip: null
|
||||
main_process_port: 29501
|
||||
main_training_function: main
|
||||
num_machines: 1
|
||||
num_processes: 8
|
||||
gpu_ids: 0,1,2,3,4,5,6,7
|
||||
use_cpu: false
|
||||
@@ -1,12 +0,0 @@
|
||||
compute_environment: LOCAL_MACHINE
|
||||
distributed_type: MULTI_GPU
|
||||
fsdp_config: {}
|
||||
machine_rank: 0
|
||||
main_process_ip: null
|
||||
main_process_port: 29501
|
||||
main_training_function: main
|
||||
mixed_precision: bf16
|
||||
num_machines: 1
|
||||
num_processes: 8
|
||||
gpu_ids: 0,1,2,3,4,5,6,7
|
||||
use_cpu: false
|
||||
@@ -1,16 +0,0 @@
|
||||
node30 slots=8
|
||||
node005 slots=8
|
||||
node006 slots=8
|
||||
node007 slots=8
|
||||
node008 slots=8
|
||||
node010 slots=8
|
||||
node012 slots=8
|
||||
node013 slots=8
|
||||
node019 slots=8
|
||||
node026 slots=8
|
||||
node027 slots=8
|
||||
node028 slots=8
|
||||
node029 slots=8
|
||||
node031 slots=8
|
||||
node059 slots=8
|
||||
node060 slots=8
|
||||
@@ -1,18 +0,0 @@
|
||||
compute_environment: LOCAL_MACHINE
|
||||
distributed_type: DEEPSPEED
|
||||
deepspeed_config:
|
||||
deepspeed_config_file: scripts/accelerate_configs/zero2.json
|
||||
deepspeed_hostfile: scripts/accelerate_configs/hostfile
|
||||
fsdp_config: {}
|
||||
machine_rank: 0
|
||||
main_process_ip: 100.64.24.30
|
||||
main_process_port: 29522
|
||||
main_training_function: main
|
||||
num_machines: 16
|
||||
num_processes: 128
|
||||
rdzv_backend: static
|
||||
same_network: true
|
||||
tpu_env: []
|
||||
tpu_use_cluster: false
|
||||
tpu_use_sudo: false
|
||||
use_cpu: false
|
||||
@@ -1,13 +0,0 @@
|
||||
compute_environment: LOCAL_MACHINE
|
||||
distributed_type: MULTI_GPU
|
||||
fsdp_config: {}
|
||||
main_process_port: 29501
|
||||
main_training_function: main
|
||||
num_machines: 2
|
||||
num_processes: 16
|
||||
rdzv_backend: static
|
||||
same_network: true
|
||||
tpu_env: []
|
||||
tpu_use_cluster: false
|
||||
tpu_use_sudo: false
|
||||
use_cpu: false
|
||||
@@ -1,13 +0,0 @@
|
||||
compute_environment: LOCAL_MACHINE
|
||||
distributed_type: DEEPSPEED
|
||||
deepspeed_config:
|
||||
deepspeed_config_file: scripts/accelerate_configs/zero2.json
|
||||
fsdp_config: {}
|
||||
machine_rank: 0
|
||||
main_process_ip: null
|
||||
main_process_port: 29503
|
||||
main_training_function: main
|
||||
num_machines: 1
|
||||
num_processes: 1
|
||||
# gpu_ids: 0
|
||||
use_cpu: false
|
||||
@@ -1,25 +0,0 @@
|
||||
{
|
||||
"fp16": {
|
||||
"enabled": false,
|
||||
"loss_scale": 0,
|
||||
"loss_scale_window": 1000,
|
||||
"initial_scale_power": 16,
|
||||
"hysteresis": 2,
|
||||
"min_loss_scale": 1
|
||||
},
|
||||
"bf16": {
|
||||
"enabled": "auto"
|
||||
},
|
||||
"communication_data_type": "fp32",
|
||||
"gradient_clipping": 1.0,
|
||||
"train_micro_batch_size_per_gpu": "auto",
|
||||
"train_batch_size": "auto",
|
||||
"gradient_accumulation_steps": "auto",
|
||||
"zero_optimization": {
|
||||
"stage": 2,
|
||||
"overlap_comm": true,
|
||||
"contiguous_gradients": true,
|
||||
"sub_group_size": 1e9,
|
||||
"reduce_bucket_size": 5e8
|
||||
}
|
||||
}
|
||||
@@ -1,26 +0,0 @@
|
||||
{
|
||||
"fp16": {
|
||||
"enabled": "auto",
|
||||
"loss_scale": 0,
|
||||
"loss_scale_window": 1000,
|
||||
"initial_scale_power": 16,
|
||||
"hysteresis": 2,
|
||||
"min_loss_scale": 1
|
||||
},
|
||||
"bf16": {
|
||||
"enabled": "auto"
|
||||
},
|
||||
"train_micro_batch_size_per_gpu": "auto",
|
||||
"train_batch_size": "auto",
|
||||
"gradient_accumulation_steps": "auto",
|
||||
"zero_optimization": {
|
||||
"stage": 2,
|
||||
"offload_optimizer": {
|
||||
"device": "cpu"
|
||||
},
|
||||
"overlap_comm": true,
|
||||
"contiguous_gradients": true,
|
||||
"sub_group_size": 1e9,
|
||||
"reduce_bucket_size": 5e8
|
||||
}
|
||||
}
|
||||
@@ -1,28 +0,0 @@
|
||||
{
|
||||
"fp16": {
|
||||
"enabled": "auto",
|
||||
"loss_scale": 0,
|
||||
"loss_scale_window": 1000,
|
||||
"initial_scale_power": 16,
|
||||
"hysteresis": 2,
|
||||
"min_loss_scale": 1
|
||||
},
|
||||
"bf16": {
|
||||
"enabled": "auto"
|
||||
},
|
||||
"train_micro_batch_size_per_gpu": "auto",
|
||||
"train_batch_size": "auto",
|
||||
"gradient_accumulation_steps": "auto",
|
||||
"zero_optimization": {
|
||||
"stage": 3,
|
||||
"overlap_comm": true,
|
||||
"contiguous_gradients": true,
|
||||
"sub_group_size": 1e9,
|
||||
"reduce_bucket_size": 5e8,
|
||||
"stage3_prefetch_bucket_size": 5e8,
|
||||
"stage3_param_persistence_threshold": 1e6,
|
||||
"stage3_max_live_parameters": 1e9,
|
||||
"stage3_max_reuse_distance": 1e9,
|
||||
"stage3_gather_16bit_weights_on_model_save": true
|
||||
}
|
||||
}
|
||||
@@ -1,39 +0,0 @@
|
||||
{
|
||||
"fp16": {
|
||||
"enabled": "auto",
|
||||
"loss_scale": 0,
|
||||
"loss_scale_window": 1000,
|
||||
"initial_scale_power": 16,
|
||||
"hysteresis": 2,
|
||||
"min_loss_scale": 1
|
||||
},
|
||||
"bf16": {
|
||||
"enabled": "auto"
|
||||
},
|
||||
"zero_optimization": {
|
||||
"stage": 3,
|
||||
"offload_optimizer": {
|
||||
"device": "cpu",
|
||||
"pin_memory": true
|
||||
},
|
||||
"offload_param": {
|
||||
"device": "cpu",
|
||||
"pin_memory": true
|
||||
},
|
||||
"overlap_comm": true,
|
||||
"contiguous_gradients": true,
|
||||
"sub_group_size": 1e9,
|
||||
"reduce_bucket_size": 5e8,
|
||||
"stage3_prefetch_bucket_size": "auto",
|
||||
"stage3_param_persistence_threshold": "auto",
|
||||
"stage3_max_live_parameters": 1e9,
|
||||
"stage3_max_reuse_distance": 1e9,
|
||||
"gather_16bit_weights_on_model_save": true
|
||||
},
|
||||
"gradient_accumulation_steps": "auto",
|
||||
"gradient_clipping": "auto",
|
||||
"train_batch_size": "auto",
|
||||
"train_micro_batch_size_per_gpu": "auto",
|
||||
"steps_per_print": 1e5,
|
||||
"wall_clock_breakdown": false
|
||||
}
|
||||
+17
-76
@@ -1,53 +1,23 @@
|
||||
|
||||
accelerate launch \
|
||||
--config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \
|
||||
torchrun --nnodes 1 --nproc_per_node 4 \
|
||||
fastvideo/train.py \
|
||||
--seed 42 \
|
||||
--pretrained_model_name_or_path data/mochi \
|
||||
--cache_dir "data/.cache" \
|
||||
--data_json_path "data/Encoder_Overfit_Data/videos2caption.json" \
|
||||
--validation_prompt_dir "data/Encoder_Overfit_Data/validation_prompt_embed_mask" \
|
||||
--uncond_prompt_dir "data/Encoder_Overfit_Data/uncond_prompt_embed_mask" \
|
||||
--gradient_checkpointing \
|
||||
--train_batch_size=1 \
|
||||
--num_latent_t 14 \
|
||||
--sp_size 2 \
|
||||
--train_sp_batch_size 2 \
|
||||
--num_latent_t 1 \
|
||||
--sp_size 1 \
|
||||
--train_sp_batch_size 1 \
|
||||
--dataloader_num_workers 1 \
|
||||
--gradient_accumulation_steps=2 \
|
||||
--gradient_accumulation_steps=1 \
|
||||
--max_train_steps=2000 \
|
||||
--learning_rate=1e-5 \
|
||||
--mixed_precision="bf16" \
|
||||
--weighting_scheme "uniform" \
|
||||
--checkpointing_steps=500 \
|
||||
--checkpoints_total_limit 3 \
|
||||
--allow_tf32 \
|
||||
--ema_start_step 0 \
|
||||
--cfg 0.1 \
|
||||
--ema_decay 0.999 \
|
||||
--output_dir="data/outputs/debug_uniform_14_sp2_spb2_4videos_2K"
|
||||
|
||||
|
||||
|
||||
|
||||
accelerate launch \
|
||||
--config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \
|
||||
fastvideo/train.py \
|
||||
--seed 42 \
|
||||
--pretrained_model_name_or_path data/mochi \
|
||||
--cache_dir "data/.cache" \
|
||||
--data_json_path "data/Encoder_Overfit_Data/videos2caption.json" \
|
||||
--gradient_checkpointing \
|
||||
--train_batch_size=1 \
|
||||
--num_latent_t 20 \
|
||||
--sp_size 4 \
|
||||
--train_sp_batch_size 1 \
|
||||
--dataloader_num_workers 1 \
|
||||
--gradient_accumulation_steps=2 \
|
||||
--max_train_steps=500 \
|
||||
--learning_rate=0.0 \
|
||||
--mixed_precision="bf16" \
|
||||
--weighting_scheme "uniform" \
|
||||
--checkpointing_steps=100 \
|
||||
--validation_steps 2 \
|
||||
--checkpointing_steps=200 \
|
||||
--validation_steps 200 \
|
||||
--validation_sampling_steps 64 \
|
||||
--checkpoints_total_limit 3 \
|
||||
--allow_tf32 \
|
||||
@@ -55,48 +25,19 @@ accelerate launch \
|
||||
--cfg 0.1 \
|
||||
--ema_decay 0.999 \
|
||||
--log_validation \
|
||||
--validation_prompt_dir "data/Encoder_Overfit_Data/validation_prompt_embed_mask" \
|
||||
--uncond_prompt_dir "data/Encoder_Overfit_Data/uncond_prompt_embed_mask" \
|
||||
--output_dir="data/outputs/debug_logvalidation_sp2"
|
||||
--output_dir="data/outputs/BW_Testrun" \
|
||||
--use_cpu_offload
|
||||
|
||||
|
||||
|
||||
torchrun --nnodes=1 --nproc_per_node=4 --master_port 29503 debug_mochi_sp.py
|
||||
torchrun --nnodes=1 --nproc_per_node=4 --master_port 29503 debug_OSP_A2A.py
|
||||
|
||||
|
||||
for seed in 42
|
||||
do
|
||||
torchrun --nnodes=1 --nproc_per_node=4 --master_port 29503 fastvideo/model/test.py --seed $seed
|
||||
done
|
||||
|
||||
# SP single node: 13
|
||||
# SP multi node: 20
|
||||
# SP multi node with larger batch size: 15
|
||||
# No SP Single node: 10.5
|
||||
# No SP Multe node: 11
|
||||
|
||||
|
||||
num_gpus=1
|
||||
|
||||
torchrun --nproc_per_node=$num_gpus fastvideo/sample/generate_synthetic.py \
|
||||
--model_path data/mochi \
|
||||
--num_frames 163 \
|
||||
--height 480 \
|
||||
--width 848 \
|
||||
--num_inference_steps 2 \
|
||||
--guidance_scale 4.5 \
|
||||
--prompt_path "data/prompt.txt" \
|
||||
--dataset_output_dir data/synthetic_debug3
|
||||
|
||||
|
||||
|
||||
|
||||
"A person is sawing a piece of wood placed on a support, using a large hand saw.
|
||||
The saw moves rapidly back and forth across the top of the wood, generating small wood chips and dust.
|
||||
The person's hands are visible, gripping the saw firmly. The background shows a workshop environment with indistinct objects and a bench.
|
||||
The focus remains on the action of sawing, with no significant change in the camera angle or background.
|
||||
The sawing action continues with the saw moving through the wood, which shows a visible cut deepening as the saw progresses.
|
||||
Wood chips continue to scatter from the cutting area.
|
||||
The person's grip adjusts slightly for better control and force.
|
||||
The background remains consistent with the workshop setting, and the camera stays focused on the sawing action without any shifts in perspective
|
||||
No new objects or persons enter the scene.
|
||||
The sawing action reaches a deeper cut into the wood, indicating progress in the task.
|
||||
The person's movements become more pronounced as they apply more force.
|
||||
The saw blade appears more prominently as it moves back and forth.
|
||||
The background and setting remain unchanged, maintaining the focus on the sawing activity.
|
||||
The camera continues to capture the action closely, emphasizing the interaction between the saw and the wood."
|
||||
|
||||
Reference in New Issue
Block a user