diff --git a/fastvideo/sample/generate_synthetic.py b/fastvideo/sample/generate_synthetic.py index 076e0efb5..f53c6c05e 100644 --- a/fastvideo/sample/generate_synthetic.py +++ b/fastvideo/sample/generate_synthetic.py @@ -19,7 +19,7 @@ def generate_video_and_latent(pipe, prompt, height, width, num_frames, num_infer generator=generator, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, - output_type="latent_and_video" + return_all_states=True, ) return video[0], latent[0], prompt_embed[0], prompt_attention_mask[0] diff --git a/fastvideo/sample/monkey_patch/sp_attention.py b/fastvideo/sample/monkey_patch/sp_attention.py new file mode 100644 index 000000000..db999d594 --- /dev/null +++ b/fastvideo/sample/monkey_patch/sp_attention.py @@ -0,0 +1,146 @@ +import inspect +import math +from typing import Callable, List, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +from torch import nn +from einops import rearrange, repeat + +from diffusers.image_processor import IPAdapterMaskProcessor +from diffusers.utils import deprecate, logging +from diffusers.utils.import_utils import is_torch_npu_available, is_xformers_available +from diffusers.utils.torch_utils import is_torch_version, maybe_allow_in_graph +from diffusers.models.embeddings import apply_rotary_emb + +import diffusers +from types import MethodType + +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_to_all_BHSD, all_gather_BHSD + +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: + + 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, *image_rotary_emb) + key = apply_rotary_emb(key, *image_rotary_emb) + + 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) + total_sequence_length = sequence_length * nccl_info.world_size + encoder_sequence_length = encoder_query.size(2) + print("sequence_length: ", sequence_length) + print("encoder_sequence_length: ", encoder_sequence_length) + + # if sp query slice, encoder_query full + # first all-to-all query -> full query, then concat encoder_query + if get_sequence_parallel_state(): + query = all_to_all_BHSD(query, scatter_dim=1, gather_dim=2).contiguous() + key = all_to_all_BHSD(key, scatter_dim=1, gather_dim=2).contiguous() + value = all_to_all_BHSD(value, scatter_dim=1, gather_dim=2).contiguous() + + if get_sequence_parallel_state(): + def shrink_head(encoder_state): + world_size, rank = nccl_info.world_size, nccl_info.rank + return rearrange(encoder_state, 'b (n h) s d -> b n h s d', n=world_size, + h=encoder_state.shape[1] // world_size).contiguous()[:, rank, ...] + encoder_query = shrink_head(encoder_query) + encoder_key = shrink_head(encoder_key) + encoder_value = shrink_head(encoder_value) + + # Hint: please check encoder_query.shape + print("query.shape: ", query.shape) + print("encoder_query.shape: ", 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) + + print("hidden_states.shape: ", hidden_states.shape) + + hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False) + + print("after attn before sp.shape", hidden_states.shape) + hidden_states, encoder_hidden_states = hidden_states.split_with_sizes( + (total_sequence_length, encoder_sequence_length), dim=2 + ) + print("after attn before sp.shape", hidden_states.shape) + if get_sequence_parallel_state(): + hidden_states = all_to_all_BHSD(hidden_states, scatter_dim=2, gather_dim=1).contiguous() + encoder_hidden_states = all_gather_BHSD(encoder_hidden_states, dim=1).contiguous() + print("after attn after sp.shape", hidden_states.shape) + + 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) + + print("after trans.shape", hidden_states.shape, encoder_hidden_states.shape) + + + # 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 + diff --git a/fastvideo/sample/monkey_patch/sp_rope.py b/fastvideo/sample/monkey_patch/sp_rope.py new file mode 100644 index 000000000..a1f3a1566 --- /dev/null +++ b/fastvideo/sample/monkey_patch/sp_rope.py @@ -0,0 +1,32 @@ +import inspect +import math +from typing import Callable, List, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +from torch import nn +from einops import rearrange, repeat + +from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info + +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, device=device, dtype=dtype) + if get_sequence_parallel_state(): + rank = nccl_info.rank + t += rank * num_frames + 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 \ No newline at end of file diff --git a/fastvideo/sample/mochi_pipe.py b/fastvideo/sample/pipeline_mochi.py similarity index 97% rename from fastvideo/sample/mochi_pipe.py rename to fastvideo/sample/pipeline_mochi.py index f95221662..bdccceda6 100644 --- a/fastvideo/sample/mochi_pipe.py +++ b/fastvideo/sample/pipeline_mochi.py @@ -32,7 +32,8 @@ from diffusers.utils.torch_utils import randn_tensor from diffusers.video_processor import VideoProcessor from diffusers.pipelines.pipeline_utils import DiffusionPipeline from diffusers.pipelines.mochi.pipeline_output import MochiPipelineOutput - +from einops import rearrange +from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info if is_torch_xla_available(): import torch_xla.core.xla_model as xm @@ -492,6 +493,7 @@ class MochiPipeline(DiffusionPipeline): callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, callback_on_step_end_tensor_inputs: List[str] = ["latents"], max_sequence_length: int = 256, + return_all_states = False, ): r""" Function invoked when calling the pipeline for generation. @@ -629,7 +631,11 @@ class MochiPipeline(DiffusionPipeline): generator, latents, ) - + world_size, rank = nccl_info.world_size, nccl_info.rank + if get_sequence_parallel_state(): + latents = rearrange(latents, "b t (n s) h w -> b t n s h w", n=world_size).contiguous() + latents = latents[:, :, rank, :, :, :] + # 5. Prepare timestep # from https://github.com/genmoai/models/blob/075b6e36db58f1242921deff83a1066887b9c9e1/src/mochi_preview/infer.py#L77 threshold_noise = 0.025 @@ -693,6 +699,14 @@ class MochiPipeline(DiffusionPipeline): if XLA_AVAILABLE: xm.mark_step() + if get_sequence_parallel_state(): + latents_shape = list(latents.shape) + full_shape = [latents_shape[0] * world_size] + latents_shape[1:] + all_latents = torch.zeros(full_shape, dtype=latents.dtype, device=latents.device) + torch.distributed.all_gather_into_tensor(all_latents, latents) + latents_list = list(all_latents.chunk(world_size, dim=0)) + latents = torch.cat(latents_list, dim=2) + if output_type == "latent": video = latents else: @@ -712,14 +726,14 @@ class MochiPipeline(DiffusionPipeline): latents = latents / self.vae.config.scaling_factor video = self.vae.decode(latents, return_dict=False)[0] - # TODO: Change pil hardcode - video = self.video_processor.postprocess_video(video, output_type='pil') - + video = self.video_processor.postprocess_video(video, output_type=output_type) + + if return_all_states: + return video, latents, prompt_embeds, prompt_attention_mask # Offload all models self.maybe_free_model_hooks() - if output_type == "latent_and_video": - return video, latents, prompt_embeds, prompt_attention_mask - + if not return_dict: return (video,) - return MochiPipelineOutput(frames=video) \ No newline at end of file + + return MochiPipelineOutput(frames=video) diff --git a/fastvideo/sample/sample_t2v_mochi_sp.py b/fastvideo/sample/sample_t2v_mochi_sp.py new file mode 100644 index 000000000..668e26e78 --- /dev/null +++ b/fastvideo/sample/sample_t2v_mochi_sp.py @@ -0,0 +1,67 @@ +import torch +from fastvideo.sample.pipeline_mochi import MochiPipeline +import torch.distributed as dist + +from diffusers.utils import export_to_video, load_image, load_video +from fastvideo.utils.parallel_states import initialize_sequence_parallel_state, nccl_info +import argparse +import os +import diffusers +from fastvideo.sample.monkey_patch.sp_attention import NewMochiAttnProcessor2_0 +from fastvideo.sample.monkey_patch.sp_rope import NewMochiRoPE + +def initialize_distributed(): + local_rank = int(os.getenv('RANK', 0)) + world_size = int(os.getenv('WORLD_SIZE', 1)) + print('world_size', world_size) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=local_rank) + initialize_sequence_parallel_state(world_size) + + +def replace_hf_ckpt_with_new_ckpt(): + initialize_distributed() + diffusers.models.attention_processor.MochiAttnProcessor2_0.__call__ = NewMochiAttnProcessor2_0.__call__ + diffusers.models.transformers.transformer_mochi.MochiRoPE._get_positions = NewMochiRoPE._get_positions + +def main(args): + replace_hf_ckpt_with_new_ckpt() + print(nccl_info.world_size) + + torch.manual_seed(args.seed) + weight_dtype = torch.bfloat16 + device = torch.cuda.current_device() + + pipe = MochiPipeline.from_pretrained(args.model_path, torch_dtype=torch.bfloat16) + pipe.enable_vae_tiling() + pipe.to(device) + #pipe.enable_model_cpu_offload() + + # Generate videos from the input prompt + video = pipe( + prompt=args.prompt, + height=args.height, + width=args.width, + num_frames=args.num_frames, + num_inference_steps=args.num_inference_steps, + guidance_scale=args.guidance_scale, + ).frames[0] + + dist.barrier() + if nccl_info.rank <= 0: + export_to_video(video, args.output_path, fps=30) + +if __name__ == "__main__": + # arg parse + parser = argparse.ArgumentParser() + parser.add_argument("--prompt", type=str, default="A hand with delicate fingers picks up a bright yellow lemon from a wooden bowl filled with lemons and sprigs of mint against a peach-colored background. The hand gently tosses the lemon up and catches it, showcasing its smooth texture. A beige string bag sits beside the bowl, adding a rustic touch to the scene. Additional lemons, one halved, are scattered around the base of the bowl. The even lighting enhances the vibrant colors and creates a fresh, inviting atmosphere.") + parser.add_argument("--num_frames", type=int, default=163) + parser.add_argument("--height", type=int, default=480) + parser.add_argument("--width", type=int, default=848) + parser.add_argument("--num_inference_steps", type=int, default=64) + parser.add_argument("--guidance_scale", type=float, default=4.5) + parser.add_argument("--model_path", type=str, default="data/mochi") + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--output_path", type=str, default="./outputs.mp4") + args = parser.parse_args() + main(args) diff --git a/fastvideo/train.py b/fastvideo/train.py index 342d703cb..4321e4406 100644 --- a/fastvideo/train.py +++ b/fastvideo/train.py @@ -28,7 +28,7 @@ import accelerate import torch from torch.nn import functional as F import transformers -from accelerate import Accelerator +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 @@ -50,7 +50,7 @@ from fastvideo.dataset.latent_datasets import LatentDataset, latent_collate_func # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.24.0") +check_min_version("0.31.0") logger = get_logger(__name__) @@ -73,6 +73,7 @@ class ForkedPdb(pdb.Pdb): @torch.inference_mode() def log_validation(args, transformer,vae, accelerator, weight_dtype, global_step, ema=False): + #TODO logger.info(f"Running validation....\n") transformer = accelerator.unwrap_model(transformer) mochi_pipeline = MochiPipeline.from_pretrained(args.pretrained_model_name_or_path,vae=vae, transformer=transformer, torch_dtype=weight_dtype).to(device=accelerator.device) @@ -229,30 +230,41 @@ def main(args): def load_model_hook(models, input_dir): - # TODO 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 - for i in range(len(models)): - # pop models so that they are not loaded again - model = models.pop() + if not accelerator.distributed_type == DistributedType.DEEPSPEED: + while len(models) > 0: + model = models.pop() - # load diffusers style into model - load_model = MochiTransformer3DModel.from_pretrained(input_dir, subfolder="transformer") - model.register_to_config(**load_model.config) - - model.load_state_dict(load_model.state_dict()) - del load_model + 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: + if args.allow_tf32 and torch.cuda.is_available(): torch.backends.cuda.matmul.allow_tf32 = True if args.scale_lr: @@ -267,7 +279,6 @@ def main(args): params_to_optimize = transformer.parameters() - # TODO: Other optmizer in cogvideoX-factory optimizer = torch.optim.AdamW( params_to_optimize, lr=args.learning_rate, diff --git a/fastvideo/utils/communications.py b/fastvideo/utils/communications.py index 923c9467b..94745a5a0 100644 --- a/fastvideo/utils/communications.py +++ b/fastvideo/utils/communications.py @@ -85,10 +85,10 @@ class _AllToAll(torch.autograd.Function): None, ) -def all_to_all_SBH( +def all_to_all_BHSD( input_: torch.Tensor, scatter_dim: int = 1, - gather_dim: int = 0, + gather_dim: int = 2, ): return _AllToAll.apply(input_, scatter_dim, gather_dim, _single_all_to_all) @@ -99,6 +99,55 @@ def all_to_all_BSND( ): return _AllToAll.apply(input_, scatter_dim, gather_dim, _all_to_all) +class _AllGather(torch.autograd.Function): + """All-gather communication with autograd support. + + Args: + input_: input tensor + dim: dimension along which to concatenate + """ + + @staticmethod + def forward(ctx, input_, dim): + ctx.dim = dim + world_size = nccl_info.world_size + group = nccl_info.group + input_size = list(input_.size()) + + ctx.input_size = input_size[dim] + + tensor_list = [torch.empty_like(input_) for _ in range(world_size)] + input_ = input_.contiguous() + dist.all_gather(tensor_list, input_, group=group) + + output = torch.cat(tensor_list, dim=dim) + return output + + @staticmethod + def backward(ctx, grad_output): + world_size = nccl_info.world_size + rank = nccl_info.rank + dim = ctx.dim + input_size = ctx.input_size + + sizes = [input_size] * world_size + + grad_input_list = torch.split(grad_output, sizes, dim=dim) + grad_input = grad_input_list[rank] + + return grad_input, None + +def all_gather_BHSD(input_: torch.Tensor, dim: int = 1): + """Performs an all-gather operation on the input tensor along the specified dimension. + + Args: + input_ (torch.Tensor): Input tensor of shape [B, H, S, D]. + dim (int, optional): Dimension along which to concatenate. Defaults to 1. + + Returns: + torch.Tensor: Output tensor after all-gather operation, concatenated along 'dim'. + """ + return _AllGather.apply(input_, dim) def prepare_parallel_data(hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask): def all_to_all(hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask): diff --git a/fastvideo/utils/parallel_states.py b/fastvideo/utils/parallel_states.py index 7b7465817..2cedb444c 100644 --- a/fastvideo/utils/parallel_states.py +++ b/fastvideo/utils/parallel_states.py @@ -28,7 +28,6 @@ def initialize_sequence_parallel_group(sequence_parallel_size): rank = int(os.getenv('RANK', '0')) world_size = int(os.getenv("WORLD_SIZE", '1')) assert world_size % sequence_parallel_size == 0, "world_size must be divisible by sequence_parallel_size, but got world_size: {}, sequence_parallel_size: {}".format(world_size, sequence_parallel_size) - # hccl nccl_info.world_size = sequence_parallel_size nccl_info.rank = rank num_sequence_parallel_groups: int = world_size // sequence_parallel_size diff --git a/infenrece_mochi.sh b/infenrece_mochi.sh index a1fa85754..38a9cf258 100644 --- a/infenrece_mochi.sh +++ b/infenrece_mochi.sh @@ -8,4 +8,4 @@ python fastvideo/sample/sample_t2v_mochi.py \ --guidance_scale 4.5 \ --seed 12346 \ --transformer_path data/outputs/debug_1e-5/checkpoint-150/transformer \ - --output_path step150__79.mp4 \ No newline at end of file + --output_path step150__79.mp4 diff --git a/infenrece_mochi_sp.sh b/infenrece_mochi_sp.sh new file mode 100644 index 000000000..1bb1f3a85 --- /dev/null +++ b/infenrece_mochi_sp.sh @@ -0,0 +1,14 @@ +#!/bin/bash + +num_gpus=2 + +torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \ + fastvideo/sample/sample_t2v_mochi_sp.py \ + --model_path data/mochi \ + --prompt "A hand with delicate fingers picks up a bright yellow lemon from a wooden bowl filled with lemons and sprigs of mint against a peach-colored background. The hand gently tosses the lemon up and catches it, showcasing its smooth texture. A beige string bag sits beside the bowl, adding a rustic touch to the scene. Additional lemons, one halved, are scattered around the base of the bowl. The even lighting enhances the vibrant colors and creates a fresh, inviting atmosphere." \ + --num_frames 163 \ + --height 480 \ + --width 848 \ + --num_inference_steps 10 \ + --guidance_scale 4.5 \ + --seed 42 \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index f73a84a70..8c5ad762f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,7 +21,7 @@ dependencies = [ "timm==1.0.11", "torchdiffeq==0.2.4", "torchmetrics==1.5.1", "tqdm==4.66.5", "urllib3==2.2.0", "uvicorn==0.32.0", "scikit-video==1.1.11", "imageio-ffmpeg==0.5.1", "sentencepiece==0.2.0", "beautifulsoup4==4.12.3", "ftfy==6.3.0", "moviepy==1.0.3", "wandb==0.18.5", "tensorboard==2.18.0", "pydantic==2.9.2", "gradio==5.3.0", "huggingface_hub==0.26.1", "protobuf==5.28.3", -"watch", "gpustat" +"watch", "gpustat", "bitsandbytes", "prodigyopt" ] [project.optional-dependencies] diff --git a/t2v_debug_multi.sh b/t2v_debug_multi.sh index 92dcb0a9a..af35d83e2 100644 --- a/t2v_debug_multi.sh +++ b/t2v_debug_multi.sh @@ -9,13 +9,14 @@ accelerate launch \ --train_batch_size=1 \ --num_latent_t 14 \ --dataloader_num_workers 1 \ - --gradient_accumulation_steps=4 \ + --gradient_accumulation_steps=1 \ --max_train_steps=200 \ --learning_rate=1e-6 \ --lr_scheduler="constant" \ --lr_warmup_steps=0 \ + --optimizer="AdamW" \ --mixed_precision="bf16" \ - --checkpointing_steps=50 \ + --checkpointing_steps=1 \ --allow_tf32 \ --tile_sample_stride 192 \ --ema_start_step 0 \ @@ -24,4 +25,4 @@ accelerate launch \ --enable_tiling \ --sp_size 1 \ --train_sp_batch_size 1 \ - --output_dir="data/outputs/debug_1e-6" + --output_dir="data/outputs/debug" diff --git a/t2v_debug_single.sh b/t2v_debug_single.sh index bf1eda2b4..db479429d 100644 --- a/t2v_debug_single.sh +++ b/t2v_debug_single.sh @@ -8,12 +8,14 @@ accelerate launch \ --data_merge_path "data/Mochi-Synthetic-Data/merge.txt" \ --gradient_checkpointing \ --train_batch_size=1 \ + --num_latent_t 14 \ --dataloader_num_workers 1 \ --gradient_accumulation_steps=1 \ --max_train_steps=1000000 \ --learning_rate=1e-4 \ --lr_scheduler="constant" \ --lr_warmup_steps=0 \ + --optimizer="AdamW" \ --mixed_precision="bf16" \ --checkpointing_steps=250 \ --allow_tf32 \