diff --git a/dfm/src/common/tokenizers/__init__ .py b/dfm/src/common/tokenizers/__init__ .py new file mode 100644 index 00000000..e69de29b diff --git a/dfm/src/common/tokenizers/cosmos/__init__.py b/dfm/src/common/tokenizers/cosmos/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/dfm/src/common/tokenizers/cosmos/cosmos1/__init__.py b/dfm/src/common/tokenizers/cosmos/cosmos1/__init__.py new file mode 100644 index 00000000..dac9a4d7 --- /dev/null +++ b/dfm/src/common/tokenizers/cosmos/cosmos1/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. diff --git a/dfm/src/common/tokenizers/cosmos/cosmos1/causal_video_tokenizer.py b/dfm/src/common/tokenizers/cosmos/cosmos1/causal_video_tokenizer.py new file mode 100644 index 00000000..0bf57f5b --- /dev/null +++ b/dfm/src/common/tokenizers/cosmos/cosmos1/causal_video_tokenizer.py @@ -0,0 +1,193 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +# pylint: disable=C0115,C0116,C0301 + +"""A library for Causal Video Tokenizer inference.""" + +from pathlib import Path + +import numpy as np +import torch +from huggingface_hub import get_token as get_hf_token +from huggingface_hub import hf_hub_download +from tqdm import tqdm + +from dfm.src.common.tokenizers.cosmos.cosmos1.video_tokenizer_utils import ( + load_jit_model, + numpy2tensor, + pad_video_batch, + tensor2numpy, + unpad_video_batch, +) + + +class CausalVideoTokenizer(torch.nn.Module): + def __init__( + self, + checkpoint_dir: str = None, + load_full_model: bool = True, + load_dec_model: bool = True, + load_enc_model: bool = True, + device: str = "cuda", + dtype: str = "bfloat16", + ) -> None: + super().__init__() + + checkpoint = Path(checkpoint_dir) + self._full_model_path = str(checkpoint / "autoencoder.jit") + self._enc_model_path = str(checkpoint / "encoder.jit") + self._dec_model_path = str(checkpoint / "decoder.jit") + self._dtype = dtype + self._device = device + + self._full_model = load_jit_model(self._full_model_path, self._device) if load_full_model else None + self._enc_model = load_jit_model(self._enc_model_path, self._device) if load_enc_model else None + self._dec_model = load_jit_model(self._dec_model_path, self._device) if load_dec_model else None + + @classmethod + def from_pretrained( + cls, + tokenizer_type="Cosmos-Tokenizer-DV4x8x8", + load_encoder=True, + load_decoder=True, + load_full_model=False, + use_pytorch=False, + dtype="bfloat16", + ): + cls._hf_model_name = f"nvidia/{tokenizer_type}" + + # Requires setting HF_TOKEN env variable + hf_token = get_hf_token() + + full_model_path = hf_hub_download( + repo_id=cls._hf_model_name, + filename="autoencoder.jit", + token=hf_token, + ) + + _ = hf_hub_download( + repo_id=cls._hf_model_name, + filename="encoder.jit", + token=hf_token, + ) + + _ = hf_hub_download( + repo_id=cls._hf_model_name, + filename="decoder.jit", + token=hf_token, + ) + + # No need to load in encoder and decoder with full model loaded + if load_full_model: + load_encoder = False + load_decoder = False + + # Assumes HF downloads all files to same local dir + ckpt_dir = str(Path(full_model_path).parent) + args = { + "checkpoint_dir": ckpt_dir, + "dtype": dtype, + "load_enc_model": load_encoder, + "load_dec_model": load_decoder, + "load_full_model": load_full_model, + } + return cls(**args) + + @torch.no_grad() + def autoencode(self, input_tensor: torch.Tensor) -> torch.Tensor: + """Reconstrcuts a batch of video tensors after embedding into a latent. + + Args: + video: The input video Bx3xTxHxW layout, range [-1..1]. + Returns: + The reconstructed video, layout Bx3xTxHxW, range [-1..1]. + """ + if self._full_model is not None: + output_tensor = self._full_model(input_tensor) + output_tensor = output_tensor[0] if isinstance(output_tensor, tuple) else output_tensor + else: + output_latent = self.encode(input_tensor)[0] + output_tensor = self.decode(output_latent) + return output_tensor + + @torch.no_grad() + def encode(self, input_tensor: torch.Tensor) -> tuple[torch.Tensor]: + """Encodes a numpy video into a CausalVideo latent or code. + + Args: + input_tensor: The input tensor Bx3xTxHxW layout, range [-1..1]. + Returns: + For causal continuous video (CV) tokenizer, the tuple contains: + - The latent embedding, Bx16x(t)x(h)x(w), where the compression + rate is (T/t x H/h x W/w), and channel dimension of 16. + For causal discrete video (DV) tokenizer, the tuple contains: + 1) The indices, Bx(t)x(h)x(w), from a codebook of size 64K, which + is formed by FSQ levels of (8,8,8,5,5,5). + 2) The discrete code, Bx6x(t)x(h)x(w), where the compression rate + is again (T/t x H/h x W/w), and channel dimension of 6. + """ + assert input_tensor.ndim == 5, "input video should be of 5D." + + output_latent = self._enc_model(input_tensor) + if isinstance(output_latent, torch.Tensor): + return output_latent + return output_latent[:-1] + + @torch.no_grad() + def decode(self, input_latent: torch.Tensor) -> torch.Tensor: + """Encodes a numpy video into a CausalVideo latent. + + Args: + input_latent: The continuous latent Bx16xtxhxw for CV, + or the discrete indices Bxtxhxw for DV. + Returns: + The reconstructed tensor, layout [B,3,1+(T-1)*8,H*16,W*16] in range [-1..1]. + """ + assert input_latent.ndim >= 4, "input latent should be of 5D for continuous and 4D for discrete." + return self._dec_model(input_latent) + + def forward( + self, + video: np.ndarray, + temporal_window: int = 17, + ) -> np.ndarray: + """Reconstructs video using a pre-trained CausalTokenizer autoencoder. + Given a video of arbitrary length, the forward invokes the CausalVideoTokenizer + in a sliding manner with a `temporal_window` size. + + Args: + video: The input video BxTxHxWx3 layout, range [0..255]. + temporal_window: The length of the temporal window to process, default=25. + Returns: + The reconstructed video in range [0..255], layout BxTxHxWx3. + """ + assert video.ndim == 5, "input video should be of 5D." + num_frames = video.shape[1] # can be of any length. + output_video_list = [] + for idx in tqdm(range(0, (num_frames - 1) // temporal_window + 1)): + # Input video for the current window. + start, end = idx * temporal_window, (idx + 1) * temporal_window + input_video = video[:, start:end, ...] + + # Spatio-temporally pad input_video so it's evenly divisible. + padded_input_video, crop_region = pad_video_batch(input_video) + input_tensor = numpy2tensor(padded_input_video, dtype=self._dtype, device=self._device) + output_tensor = self.autoencode(input_tensor) + padded_output_video = tensor2numpy(output_tensor) + output_video = unpad_video_batch(padded_output_video, crop_region) + + output_video_list.append(output_video) + return np.concatenate(output_video_list, axis=1) diff --git a/dfm/src/common/tokenizers/cosmos/cosmos1/video_tokenizer_utils.py b/dfm/src/common/tokenizers/cosmos/cosmos1/video_tokenizer_utils.py new file mode 100644 index 00000000..8fb67279 --- /dev/null +++ b/dfm/src/common/tokenizers/cosmos/cosmos1/video_tokenizer_utils.py @@ -0,0 +1,145 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Utility functions for the inference libraries.""" + +import numpy as np +import torch + + +_DTYPE, _DEVICE = torch.bfloat16, "cuda" +_UINT8_MAX_F = float(torch.iinfo(torch.uint8).max) +_SPATIAL_ALIGN = 16 +_TEMPORAL_ALIGN = 8 + + +def load_jit_model(jit_filepath: str = None, device: str = "cuda") -> torch.jit.ScriptModule: + """Loads a torch.jit.ScriptModule from a filepath. + + Args: + jit_filepath: The filepath to the JIT-compiled model. + device: The device to load the model onto, default=cuda. + Returns: + The JIT compiled model loaded to device and on eval mode. + """ + model = torch.jit.load(jit_filepath, map_location=device) + return model.eval().to(device) + + +def numpy2tensor( + input_image: np.ndarray, + dtype: torch.dtype = _DTYPE, + device: str = _DEVICE, + range_min: int = -1, +) -> torch.Tensor: + """Converts image(dtype=np.uint8) to `dtype` in range [0..255]. + + Args: + input_image: A batch of images in range [0..255], BxHxWx3 layout. + Returns: + A torch.Tensor of layout Bx3xHxW in range [-1..1], dtype. + """ + ndim = input_image.ndim + indices = list(range(1, ndim))[-1:] + list(range(1, ndim))[:-1] + image = input_image.transpose((0,) + tuple(indices)) / _UINT8_MAX_F + if range_min == -1: + image = 2.0 * image - 1.0 + return torch.from_numpy(image).to(dtype).to(device) + + +def tensor2numpy(input_tensor: torch.Tensor, range_min: int = -1) -> np.ndarray: + """Converts tensor in [-1,1] to image(dtype=np.uint8) in range [0..255]. + + Args: + input_tensor: Input image tensor of Bx3xHxW layout, range [-1..1]. + Returns: + A numpy image of layout BxHxWx3, range [0..255], uint8 dtype. + """ + if range_min == -1: + input_tensor = (input_tensor.float() + 1.0) / 2.0 + ndim = input_tensor.ndim + output_image = input_tensor.clamp(0, 1).cpu().numpy() + output_image = output_image.transpose((0,) + tuple(range(2, ndim)) + (1,)) + return (output_image * _UINT8_MAX_F + 0.5).astype(np.uint8) + + +def pad_video_batch( + batch: np.ndarray, + temporal_align: int = _TEMPORAL_ALIGN, + spatial_align: int = _SPATIAL_ALIGN, +) -> tuple[np.ndarray, list[int]]: + """Pads a batch of videos to be divisible by `temporal_align` or `spatial_align`. + + Zero pad spatially. Reflection pad temporally to handle causality better. + Args: + batch: The batch of videos to pad., layout BxFxHxWx3, in any range. + align: The alignment to pad to. + Returns: + The padded batch and the crop region. + """ + num_frames, height, width = batch.shape[-4:-1] + align = spatial_align + height_to_pad = (align - height % align) if height % align != 0 else 0 + width_to_pad = (align - width % align) if width % align != 0 else 0 + + align = temporal_align + frames_to_pad = (align - (num_frames - 1) % align) if (num_frames - 1) % align != 0 else 0 + + crop_region = [ + frames_to_pad >> 1, + height_to_pad >> 1, + width_to_pad >> 1, + num_frames + (frames_to_pad >> 1), + height + (height_to_pad >> 1), + width + (width_to_pad >> 1), + ] + batch = np.pad( + batch, + ( + (0, 0), + (0, 0), + (height_to_pad >> 1, height_to_pad - (height_to_pad >> 1)), + (width_to_pad >> 1, width_to_pad - (width_to_pad >> 1)), + (0, 0), + ), + mode="constant", + ) + batch = np.pad( + batch, + ( + (0, 0), + (frames_to_pad >> 1, frames_to_pad - (frames_to_pad >> 1)), + (0, 0), + (0, 0), + (0, 0), + ), + mode="edge", + ) + return batch, crop_region + + +def unpad_video_batch(batch: np.ndarray, crop_region: list[int]) -> np.ndarray: + """Unpads video with `crop_region`. + + Args: + batch: A batch of numpy videos, layout BxFxHxWxC. + crop_region: [f1,y1,x1,f2,y2,x2] first, top, left, last, bot, right crop indices. + + Returns: + np.ndarray: Cropped numpy video, layout BxFxHxWxC. + """ + assert len(crop_region) == 6, "crop_region should be len of 6." + f1, y1, x1, f2, y2, x2 = crop_region + return batch[..., f1:f2, y1:y2, x1:x2, :] diff --git a/dfm/src/common/utils/__init__.py b/dfm/src/common/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/dfm/src/common/utils/batch_ops.py b/dfm/src/common/utils/batch_ops.py new file mode 100644 index 00000000..956dfbee --- /dev/null +++ b/dfm/src/common/utils/batch_ops.py @@ -0,0 +1,104 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. 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 torch import Tensor + + +def common_broadcast(x: Tensor, y: Tensor) -> tuple[Tensor, Tensor]: + """ + Broadcasts two tensors to have the same shape by adding singleton dimensions where necessary. + + Args: + x (Tensor): The first input tensor. + y (Tensor): The second input tensor. + + Returns: + tuple[Tensor, Tensor]: A tuple containing the two tensors with broadcasted shapes. + + Raises: + AssertionError: If the dimensions of the tensors do not match at any axis within their common dimensions. + """ + ndims1 = x.ndim + ndims2 = y.ndim + + common_ndims = min(ndims1, ndims2) + for axis in range(common_ndims): + assert x.shape[axis] == y.shape[axis], "Dimensions not equal at axis {}".format(axis) + + if ndims1 < ndims2: + x = x.reshape(x.shape + (1,) * (ndims2 - ndims1)) + elif ndims2 < ndims1: + y = y.reshape(y.shape + (1,) * (ndims1 - ndims2)) + + return x, y + + +def batch_add(x: Tensor, y: Tensor) -> Tensor: + """ + Adds two tensors element-wise after broadcasting them to a common shape. + + Args: + x (Tensor): The first input tensor. + y (Tensor): The second input tensor. + + Returns: + Tensor: The element-wise sum of the input tensors after broadcasting. + """ + x, y = common_broadcast(x, y) + return x + y + + +def batch_mul(x: Tensor, y: Tensor) -> Tensor: + """ + Multiplies two tensors element-wise after broadcasting them to a common shape. + + Args: + x (Tensor): The first input tensor. + y (Tensor): The second input tensor. + + Returns: + Tensor: The element-wise product of the input tensors after broadcasting. + """ + x, y = common_broadcast(x, y) + return x * y + + +def batch_sub(x: Tensor, y: Tensor) -> Tensor: + """ + Subtracts two tensors element-wise after broadcasting them to a common shape. + + Args: + x (Tensor): The first input tensor. + y (Tensor): The second input tensor. + + Returns: + Tensor: The result of element-wise subtraction of the input tensors. + """ + x, y = common_broadcast(x, y) + return x - y + + +def batch_div(x: Tensor, y: Tensor) -> Tensor: + """ + Divides two tensors element-wise after broadcasting them to a common shape. + + Args: + x (Tensor): The first input tensor. + y (Tensor): The second input tensor. + + Returns: + Tensor: The result of element-wise division of `x` by `y` after broadcasting. + """ + x, y = common_broadcast(x, y) + return x / y diff --git a/dfm/src/common/utils/dynamic_import.py b/dfm/src/common/utils/dynamic_import.py new file mode 100644 index 00000000..dfe4c6dd --- /dev/null +++ b/dfm/src/common/utils/dynamic_import.py @@ -0,0 +1,47 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. 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. + +import importlib + + +def dynamic_import(full_path): + """ + Dynamically import a class or function from a given full path. + + :param full_path: The full path to the class or function (e.g., "package.module.ClassName") + :return: The imported class or function + :raises ImportError: If the module or attribute cannot be imported + :raises AttributeError: If the attribute does not exist in the module + """ + try: + # Split the full path into module path and attribute name + module_path, attribute_name = full_path.rsplit(".", 1) + except ValueError as e: + raise ImportError( + f"Invalid full path '{full_path}'. It should contain both module and attribute names." + ) from e + + # Import the module + try: + module = importlib.import_module(module_path) + except ImportError as e: + raise ImportError(f"Cannot import module '{module_path}'.") from e + + # Retrieve the attribute from the module + try: + attribute = getattr(module, attribute_name) + except AttributeError as e: + raise AttributeError(f"Module '{module_path}' does not have an attribute '{attribute_name}'.") from e + + return attribute diff --git a/dfm/src/common/utils/save_video.py b/dfm/src/common/utils/save_video.py new file mode 100644 index 00000000..64803706 --- /dev/null +++ b/dfm/src/common/utils/save_video.py @@ -0,0 +1,36 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. 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. + +import imageio +import numpy as np + + +def save_video( + grid: np.ndarray, + fps: int, + H: int, + W: int, + video_save_quality: int, + video_save_path: str, +): + kwargs = { + "fps": fps, + "quality": video_save_quality, + "macro_block_size": 1, + "ffmpeg_params": ["-s", f"{W}x{H}"], + "output_params": ["-f", "mp4"], + } + + print("video_save_path", video_save_path) + imageio.mimsave(video_save_path, grid, "mp4", **kwargs) diff --git a/dfm/src/common/utils/torch_split_tensor_for_cp.py b/dfm/src/common/utils/torch_split_tensor_for_cp.py new file mode 100644 index 00000000..f389b7ba --- /dev/null +++ b/dfm/src/common/utils/torch_split_tensor_for_cp.py @@ -0,0 +1,51 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. 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. + +import torch +from torch import Tensor +from torch.distributed import ProcessGroup, all_gather, get_world_size + + +def cat_outputs_cp(x: Tensor, seq_dim: int, cp_group: ProcessGroup) -> Tensor: + """ + Concatenates tensors from multiple processes along a specified dimension. + + This function gathers tensors from all processes in the given process group + and concatenates them along the specified dimension. + + Args: + x (Tensor): The input tensor to be gathered and concatenated. + seq_dim (int): The dimension along which to concatenate the gathered tensors. + cp_group (ProcessGroup): The process group containing all the processes involved in the gathering. + + Returns: + Tensor: A tensor resulting from the concatenation of tensors from all processes. + + Raises: + RuntimeError: If the gathering of tensors fails. + """ + # Number of processes in the group + world_size = get_world_size(cp_group) + + # List to hold tensors from each rank + gathered_tensors = [torch.zeros_like(x) for _ in range(world_size)] + + # Attempt to gather tensors from all ranks + try: + all_gather(gathered_tensors, x, group=cp_group) + except RuntimeError as e: + raise RuntimeError(f"Gathering failed: {e}") + + # Concatenate tensors along the specified dimension + return torch.cat(gathered_tensors, dim=seq_dim) diff --git a/dfm/src/megatron/data/__init__.py b/dfm/src/megatron/data/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/dfm/src/megatron/data/common/__init__.py b/dfm/src/megatron/data/common/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/dfm/src/megatron/data/common/base_energon_datamodule.py b/dfm/src/megatron/data/common/base_energon_datamodule.py new file mode 100644 index 00000000..0d6cb99d --- /dev/null +++ b/dfm/src/megatron/data/common/base_energon_datamodule.py @@ -0,0 +1,339 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. 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. + +import logging +from typing import Any, Dict, Literal, Optional + +from megatron.core import parallel_state +from megatron.energon import WorkerConfig, get_savable_loader, get_train_dataset + + +logger = logging.getLogger(__name__) + + +class EnergonMultiModalDataModule: + """ + A DataModule for handling multimodal datasets with images and text. + + This data module is designed to work with multimodal datasets that involve both images and text. + It provides a seamless interface to load training and validation data, manage batching, and handle + the state of the data pipeline across training epochs. The module integrates with the Megatron-Energon + framework for efficient data handling in large-scale distributed training. + + Attributes: + path (str): Path to the energon dataset. + tokenizer (Tokenizer): The tokenizer used for processing text. + image_processor (ImageProcessor): The image processor used for preprocessing images. + seq_length (int): The maximum sequence length for tokenized text. + micro_batch_size (int): The batch size for training and validation. + num_workers (int): Number of workers for data loading. + pin_memory (bool): Whether to pin memory in the DataLoader. + multimodal_sample_config (MultiModalSampleConfig): Configuration object for multimodal samples. + task_encoder (MultiModalTaskEncoder): Encoder responsible for encoding and batching samples. + init_global_step (int): The initial global step for the trainer, used for resuming training. + data_sampler (SequentialMegatronSampler): Sampler responsible for generating sequential samples. + train_dataloader_object (Optional): The DataLoader object for training data. + val_dataloader_object (Optional): The DataLoader object for validation data. + """ + + def __init__( + self, + path: str, + tokenizer, + image_processor, + seq_length: int = 2048, + micro_batch_size: int = 1, + global_batch_size: int = 1, + num_workers: int = 1, + num_val_workers: int | None = None, + pin_memory: bool = True, + shuffle_buffer_size: int = 100, + max_samples_per_sequence: int | None = None, + multimodal_sample_config: Optional[Any] = None, + task_encoder: Optional[Any] = None, + decoder_seq_length: Optional[int] = None, + packing_buffer_size: Optional[int] = None, + validation_task_encoder: Optional[Any] = None, + **kwargs, + ) -> None: + """ + Initialize the EnergonMultiModalDataModule. + + Parameters: + path (str): Path to the dataset. + tokenizer (Tokenizer): The tokenizer used for processing text. + image_processor (ImageProcessor): The image processor used for preprocessing images. + seq_length (int, optional): The maximum sequence length for tokenized text. Defaults to 2048. + micro_batch_size (int, optional): The batch size for training and validation. Defaults to 1. + num_workers (int, optional): Number of workers for data loading. Defaults to 1. + num_val_workers (int, optional): Number of workers for validation data loading. Defaults to num_workers. + pin_memory (bool, optional): Whether to pin memory in the DataLoader. Defaults to True. + multimodal_sample_config (MultiModalSampleConfig, optional): Configuration object for multimodal samples. + Defaults to MultiModalSampleConfig(). + shuffle_buffer_size (int, optional): Size of the shuffle buffer. Defaults to 100. + max_samples_per_sequence (int, optional): Maximum number of samples per sequence to load from memory. + Defaults to None (loads the whole tar file at once). + task_encoder (MultiModalTaskEncoder, optional): Encoder responsible for encoding and batching samples. + If not provided, a default (MultimodalTaskEncoder) encoder will be created. Defaults to None. + decoder_seq_length (int, optional): The max sequence length for the decoder. Used in encoder-decoder models + packing_buffer_size (int, optional): Size of the packing buffer for batched samples. Defaults to None. + validation_task_encoder (MultiModalTaskEncoder, optional): Encoder responsible for encoding + and batching samples for validation. Defaults to None and will be the same as task_encoder. + **kwargs: Additional keyword arguments. Will be passed to get_train_dataset() of Energon + """ + + super().__init__() + self.path = path + self.tokenizer = tokenizer + self.image_processor = image_processor + self.seq_length = seq_length + self.decoder_seq_length = decoder_seq_length + self.micro_batch_size = micro_batch_size + self.global_batch_size = global_batch_size + self.num_workers = num_workers + self.pin_memory = pin_memory + self.multimodal_sample_config = multimodal_sample_config + self.shuffle_buffer_size = shuffle_buffer_size + self.max_samples_per_sequence = max_samples_per_sequence + self.task_encoder = task_encoder + self.init_global_step = 0 + self.train_dataloader_object = None + self.val_dataloader_object = None + self.packing_buffer_size = packing_buffer_size + self.validation_task_encoder = validation_task_encoder or self.task_encoder + self.num_val_workers = num_val_workers or self.num_workers + self.kwargs = kwargs + + def datasets_provider(self, worker_config, split: Literal["train", "val"] = "val"): + """ + Provide the dataset for training or validation. + + This method retrieves the dataset for the specified split (either 'train' or 'val') and configures + it according to the worker configuration. + + Parameters: + worker_config: Configuration for the data loader workers. + split (Literal['train', 'val'], optional): The data split to retrieve ('train' or 'val'). Defaults to 'val'. + + Returns: + Dataset: The dataset configured for the specified split. + """ + + if split not in {"train", "val"}: + raise ValueError("Invalid value for split. Allowed values are 'train' or 'val'.") + + if split == "train": + task_encoder = self.task_encoder + else: + task_encoder = self.validation_task_encoder + + _dataset = get_train_dataset( + self.path, + batch_size=self.micro_batch_size, + task_encoder=task_encoder, + worker_config=worker_config, + packing_buffer_size=self.packing_buffer_size, + split_part=split, + shuffle_buffer_size=self.shuffle_buffer_size, + max_samples_per_sequence=self.max_samples_per_sequence, + **self.kwargs, + ) + + return _dataset + + def build(self): + return self.train_dataloader(), self.val_dataloader() + + def train_dataloader(self) -> Any: + """ + Initialize and return the training DataLoader. + + This method initializes the DataLoader for the training dataset. It uses the global step + from the trainer to configure the data sampler and ensures that the parallel state is initialized + correctly for distributed training. + + Returns: + TRAIN_DATALOADERS: The DataLoader for the training dataset. + """ + + logger.info(f"Multimodal train dataloader initializing with init_global_step {self.init_global_step}") + if self.train_dataloader_object: + return self.train_dataloader_object + if not parallel_state.is_initialized(): + logger.info( + f"Muiltimodal data loader parallel state is not initialized," + f"using default worker config with no_workers {self.num_workers}" + ) + worker_config = WorkerConfig.default_worker_config(self.num_workers) + else: + rank = parallel_state.get_data_parallel_rank() + world_size = parallel_state.get_data_parallel_world_size() + data_parallel_group = parallel_state.get_data_parallel_group() + logger.info( + f" Multimodal train dataloader initializing with" + f"rank {rank} world_size {world_size} data_parallel_group {data_parallel_group} ****** " + ) + worker_config = WorkerConfig( + rank=rank, + world_size=world_size, + num_workers=self.num_workers, + data_parallel_group=data_parallel_group, + worker_debug_path=None, + worker_log_level=0, + ) + train_dataset = self.datasets_provider(worker_config, split="train") + energon_dataloader = get_savable_loader(train_dataset, worker_config=worker_config) + self.train_dataloader_object = energon_dataloader + return self.train_dataloader_object + + def val_dataloader(self): + """ + Initialize and return the validation DataLoader. + + This method initializes the DataLoader for the validation dataset. It ensures that the parallel state + is initialized correctly for distributed training and returns a configured DataLoader object. + + Returns: + EVAL_DATALOADERS: The DataLoader for the validation dataset. + """ + if self.val_dataloader_object: + return self.val_dataloader_object + + if not parallel_state.is_initialized(): + logger.info( + f"Muiltimodal val data loader parallel state is not initialized," + f"using default worker config with no_workers {self.num_workers}" + ) + worker_config = WorkerConfig.default_worker_config(self.num_val_workers) + else: + rank = parallel_state.get_data_parallel_rank() + world_size = parallel_state.get_data_parallel_world_size() + data_parallel_group = parallel_state.get_data_parallel_group() + + logger.info(f"rank {rank} world_size {world_size} data_parallel_group {data_parallel_group}") + worker_config = WorkerConfig( + rank=rank, + world_size=world_size, + num_workers=self.num_workers, + data_parallel_group=data_parallel_group, + worker_debug_path=None, + worker_log_level=0, + ) + val_dataset = self.datasets_provider(worker_config, split="val") + energon_loader = get_savable_loader(val_dataset, worker_config=worker_config) + self.val_dataloader_object = energon_loader + return self.val_dataloader_object + + def test_dataloader(self) -> None: + """ + Return None as test dataset split does not exist. + + This method overrides the test_dataloader method and returns None since the test dataset split + is not defined or used in this module. + + Returns: + None + """ + logger.warning("Multimodal dataloader test dataset split does not exist") + return None + + def state_dict(self) -> Dict[str, Any]: + """ + Save the state of the data module. + + This method is called when saving a checkpoint. It generates and saves the state of the data module, + including the state of the dataloader and the number of consumed samples. + + Returns: + Dict[str, Any]: A dictionary containing the state of the data module. + """ + + if self.trainer: + dataloader_obj = self.trainer.train_dataloader + + state = [] + # All ranks should be zero except the dp rank. + if ( + parallel_state.get_context_parallel_rank() + or parallel_state.get_pipeline_model_parallel_rank() + or parallel_state.get_tensor_model_parallel_rank() + or parallel_state.get_expert_model_parallel_rank() + ) == 0: + # Save_state_global in energon assumes that we call it for only the first rank within each group that + # shares the same dataloader state. By making sure that current rank is the first rank in a model + # parallel group, we ensure this. + state = dataloader_obj.save_state_global(global_dst_rank=0) + + consumed_samples = self.data_sampler.compute_consumed_samples( + self.trainer.global_step - self.init_global_step + ) + + if state is None: + state = [] # Megatron core requires all the states on all the ranks to have same python + # type. Energon sends the state as a list + logger.info(f"Multimodal data loader saving dataloader state dict consumed samples {consumed_samples}") + return {"dataloader_state": state, "consumed_samples": consumed_samples} + + logger.warning("trainer object not connected to data module object returning empty state") + return {} + + def load_state_dict(self, state_dict: Dict[str, Any]) -> None: + """ + Load the state of the data module from a checkpoint. + + This method is called when loading a checkpoint. It restores the state of the data module, + including the state of the dataloader and the number of consumed samples. + + Parameters: + state_dict (Dict[str, Any]): The state dictionary containing the saved state of the data module. + """ + if not "dataloader_state" in state_dict: + logger.warning( + f"Data loader state cannot be resumed from state_dict, " + f"it does not have the required key dataloader_state. It has {state_dict.keys()}" + ) + return + + state = state_dict["dataloader_state"] + try: + if self.trainer: + self.trainer.datamodule.train_dataloader().restore_state_global(state) + logger.info("Multimodal dataloader state restored") + else: + logger.error(f"Cannot restore state from state_dict {state_dict}") + raise ValueError( + "Cannot restore state from state_dict: " + "Is the trainer object is initialized and attached to datamodule???" + ) + except Exception as e: + logger.warning( + f"Failed to dataloader restore state due to [Please ensure you are using same version " + f"of energon while saving and loading, Continuing without restoring data loader] : {e}" + ) + + try: + from megatron.core.num_microbatches_calculator import update_num_microbatches + + except (ImportError, ModuleNotFoundError): + logger.warning("Megatron num_microbatches_calculator not found, using Apex version.") + from apex.transformer.pipeline_parallel.utils import update_num_microbatches + + consumed_samples = state_dict["consumed_samples"] + self.data_sampler.init_consumed_samples = consumed_samples + self.data_sampler.prev_consumed_samples = consumed_samples + logger.info(f"Multimodal dataloader load state dict with consumed_samples {consumed_samples}") + update_num_microbatches( + consumed_samples=consumed_samples, + consistency_check=False, + ) diff --git a/dfm/src/megatron/data/common/diffusion_energon_datamodule.py b/dfm/src/megatron/data/common/diffusion_energon_datamodule.py new file mode 100644 index 00000000..e25cd944 --- /dev/null +++ b/dfm/src/megatron/data/common/diffusion_energon_datamodule.py @@ -0,0 +1,185 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. 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. + +# pylint: disable=C0115,C0116,C0301 + + +import logging +from dataclasses import dataclass +from typing import Any, Dict, Literal + +from megatron.bridge.data.utils import DatasetBuildContext, DatasetProvider +from megatron.energon import DefaultTaskEncoder, get_train_dataset +from torch import int_repr + +from dfm.src.megatron.data.common.base_energon_datamodule import EnergonMultiModalDataModule +from dfm.src.megatron.data.dit.dit_taskencoder import DiTTaskEncoder + + +@dataclass(kw_only=True) +class DiffusionDataModuleConfig(DatasetProvider): + path: str + seq_length: int + micro_batch_size: int + task_encoder_seq_length: int + packing_buffer_size: int + global_batch_size: int + num_workers: int_repr + dataloader_type: str = "external" + + def __post_init__(self): + self.dataset = DiffusionDataModule( + path=self.path, + seq_length=self.seq_length, + task_encoder=DiTTaskEncoder( + seq_length=self.task_encoder_seq_length, packing_buffer_size=self.packing_buffer_size + ), + micro_batch_size=self.micro_batch_size, + packing_buffer_size=self.packing_buffer_size, + global_batch_size=self.global_batch_size, + num_workers=self.num_workers, + ) + self.sequence_length = self.dataset.seq_length + + def build_datasets(self, context: DatasetBuildContext): + # TODO: add validation and test datasets + return self.dataset.train_dataloader(), self.dataset.train_dataloader(), self.dataset.train_dataloader() + + +class DiffusionDataModule(EnergonMultiModalDataModule): + """ + A PyTorch Lightning DataModule for handling multimodal datasets with images and text. + + This data module is designed to work with multimodal datasets that involve both images and text. + It provides a seamless interface to load training and validation data, manage batching, and handle + the state of the data pipeline across training epochs. The module integrates with the Megatron-Energon + framework for efficient data handling in large-scale distributed training. + + Attributes: + path (str): Path to the energon dataset. + tokenizer (Tokenizer): The tokenizer used for processing text. + image_processor (ImageProcessor): The image processor used for preprocessing images. + seq_length (int): The maximum sequence length for tokenized text. + micro_batch_size (int): The batch size for training and validation. + num_workers (int): Number of workers for data loading. + pin_memory (bool): Whether to pin memory in the DataLoader. + multimodal_sample_config (MultiModalSampleConfig): Configuration object for multimodal samples. + task_encoder (MultiModalTaskEncoder): Encoder responsible for encoding and batching samples. + init_global_step (int): The initial global step for the trainer, used for resuming training. + data_sampler (SequentialMegatronSampler): Sampler responsible for generating sequential samples. + train_dataloader_object (Optional): The DataLoader object for training data. + val_dataloader_object (Optional): The DataLoader object for validation data. + """ + + def __init__( + self, + path: str, + seq_length: int = 2048, + micro_batch_size: int = 1, + global_batch_size: int = 8, + num_workers: int = 1, + pin_memory: bool = True, + packing_buffer_size: int = None, + task_encoder: DefaultTaskEncoder = None, + use_train_split_for_val: bool = False, + ) -> None: + """ + Initialize the SimpleMultiModalDataModule. + + Parameters: + path (str): Path to the dataset. + tokenizer (Tokenizer): The tokenizer used for processing text. + image_processor (ImageProcessor): The image processor used for preprocessing images. + seq_length (int, optional): The maximum sequence length for tokenized text. Defaults to 2048. + micro_batch_size (int, optional): The batch size for training and validation. Defaults to 1. + num_workers (int, optional): Number of workers for data loading. Defaults to 1. + pin_memory (bool, optional): Whether to pin memory in the DataLoader. Defaults to True. + """ + + super().__init__( + path=path, + tokenizer=None, + image_processor=None, + seq_length=seq_length, + micro_batch_size=micro_batch_size, + global_batch_size=global_batch_size, + num_workers=num_workers, + packing_buffer_size=packing_buffer_size, + pin_memory=pin_memory, + task_encoder=task_encoder, + ) + self.use_train_split_for_val = use_train_split_for_val + + def datasets_provider(self, worker_config, split: Literal["train", "val"] = "val"): + """ + Provide the dataset for training or validation. + + This method retrieves the dataset for the specified split (either 'train' or 'val') and configures + it according to the worker configuration. + + Parameters: + worker_config: Configuration for the data loader workers. + split (Literal['train', 'val'], optional): The data split to retrieve ('train' or 'val'). Defaults to 'val'. + + Returns: + Dataset: The dataset configured for the specified split. + """ + if split not in {"train", "val"}: + raise ValueError("Invalid value for split. Allowed values are 'train' or 'val'.") + if self.use_train_split_for_val: + split = "train" + _dataset = get_train_dataset( + self.path, + batch_size=self.micro_batch_size, + packing_buffer_size=self.packing_buffer_size, + task_encoder=self.task_encoder, + worker_config=worker_config, + max_samples_per_sequence=None, + shuffle_buffer_size=100, + split_part=split, + batch_drop_last=True, + virtual_epoch_length=1_000_000_000, # a hack to avoid energon end of epoch warning + ) + return _dataset + + def val_dataloader(self): + """ + Configure the validation DataLoader. + + This method configures the DataLoader for validation data. + + Parameters: + worker_config: Configuration for the data loader workers. + + Returns: + DataLoader: The DataLoader for validation data. + """ + if self.use_train_split_for_val: + return self.train_dataloader() + return super().val_dataloader() + + def load_state_dict(self, state_dict: Dict[str, Any]) -> None: + """ + Load the state of the data module from a checkpoint. + + This method is called when loading a checkpoint. It restores the state of the data module, + including the state of the dataloader and the number of consumed samples. + + Parameters: + state_dict (Dict[str, Any]): The state dictionary containing the saved state of the data module. + """ + try: + super().load_state_dict(state_dict) + except Exception as e: + logging.warning(f"datamodule.load_state_dict failed {e}") diff --git a/dfm/src/megatron/data/common/diffusion_sample.py b/dfm/src/megatron/data/common/diffusion_sample.py new file mode 100644 index 00000000..3d5f2384 --- /dev/null +++ b/dfm/src/megatron/data/common/diffusion_sample.py @@ -0,0 +1,96 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. 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 dataclasses import dataclass +from typing import Any, Optional + +import torch +from megatron.energon import Sample + + +@dataclass +class DiffusionSample(Sample): + """ + Data class representing a sample for diffusion tasks. + + Attributes: + video (torch.Tensor): Video latents (C T H W). + t5_text_embeddings (torch.Tensor): Text embeddings (S D). + t5_text_mask (torch.Tensor): Mask for text embeddings. + loss_mask (torch.Tensor): Mask indicating valid positions for loss computation. + image_size (Optional[torch.Tensor]): Tensor containing image dimensions. + fps (Optional[torch.Tensor]): Frame rate of the video. + num_frames (Optional[torch.Tensor]): Number of frames in the video. + padding_mask (Optional[torch.Tensor]): Mask indicating padding positions. + seq_len_q (Optional[torch.Tensor]): Sequence length for query embeddings. + seq_len_kv (Optional[torch.Tensor]): Sequence length for key/value embeddings. + pos_ids (Optional[torch.Tensor]): Positional IDs. + latent_shape (Optional[torch.Tensor]): Shape of the latent tensor. + """ + + video: torch.Tensor # video latents (C T H W) + context_embeddings: torch.Tensor # (S D) + context_mask: torch.Tensor = None # 1 + image_size: Optional[torch.Tensor] = None + loss_mask: torch.Tensor = None + fps: Optional[torch.Tensor] = None + num_frames: Optional[torch.Tensor] = None + padding_mask: Optional[torch.Tensor] = None + seq_len_q: Optional[torch.Tensor] = None + seq_len_kv: Optional[torch.Tensor] = None + pos_ids: Optional[torch.Tensor] = None + latent_shape: Optional[torch.Tensor] = None + + def to_dict(self) -> dict: + """Converts the sample to a dictionary.""" + return dict( + video=self.video, + context_embeddings=self.context_embeddings, + context_mask=self.context_mask, + loss_mask=self.loss_mask, + image_size=self.image_size, + fps=self.fps, + num_frames=self.num_frames, + padding_mask=self.padding_mask, + seq_len_q=self.seq_len_q, + seq_len_kv=self.seq_len_kv, + pos_ids=self.pos_ids, + latent_shape=self.latent_shape, + ) + + def __add__(self, other: Any) -> int: + """Adds the sequence length of this sample with another sample or integer.""" + if isinstance(other, DiffusionSample): + # Combine the values of the two instances + return self.seq_len_q.item() + other.seq_len_q.item() + elif isinstance(other, int): + # Add an integer to the value + return self.seq_len_q.item() + other + raise NotImplementedError + + def __radd__(self, other: Any) -> int: + """Handles reverse addition for summing with integers.""" + # This is called if sum or other operations start with a non-DiffusionSample object. + # e.g., sum([DiffusionSample(1), DiffusionSample(2)]) -> the 0 + DiffusionSample(1) calls __radd__. + if isinstance(other, int): + return self.seq_len_q.item() + other + raise NotImplementedError + + def __lt__(self, other: Any) -> bool: + """Compares this sample's sequence length with another sample or integer.""" + if isinstance(other, DiffusionSample): + return self.seq_len_q.item() < other.seq_len_q.item() + elif isinstance(other, int): + return self.seq_len_q.item() < other + raise NotImplementedError diff --git a/dfm/src/megatron/data/common/diffusion_task_encoder_with_sp.py b/dfm/src/megatron/data/common/diffusion_task_encoder_with_sp.py new file mode 100644 index 00000000..80a253de --- /dev/null +++ b/dfm/src/megatron/data/common/diffusion_task_encoder_with_sp.py @@ -0,0 +1,113 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# 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. + + +import random +from abc import ABC, abstractmethod +from typing import List + +import torch +from megatron.energon import DefaultTaskEncoder +from megatron.energon.task_encoder.base import stateless +from megatron.energon.task_encoder.cooking import Cooker, basic_sample_keys + +from dfm.src.megatron.data.common.diffusion_sample import DiffusionSample +from dfm.src.megatron.data.common.sequence_packing_utils import first_fit_decreasing + + +def cook(sample: dict) -> dict: + """ + Processes a raw sample dictionary from energon dataset and returns a new dictionary with specific keys. + + Args: + sample (dict): The input dictionary containing the raw sample data. + + Returns: + dict: A new dictionary containing the processed sample data with the following keys: + - All keys from the result of `basic_sample_keys(sample)` + - 'json': The contains meta data like resolution, aspect ratio, fps, etc. + - 'pth': contains video latent tensor + - 'pickle': contains text embeddings + """ + return dict( + **basic_sample_keys(sample), + json=sample[".json"], + pth=sample[".pth"], + pickle=sample[".pickle"], + ) + + +class DiffusionTaskEncoderWithSequencePacking(DefaultTaskEncoder, ABC): + cookers = [ + Cooker(cook), + ] + + def __init__( + self, + *args, + max_frames: int = None, + text_embedding_padding_size: int = 512, + seq_length: int = None, + patch_spatial: int = 2, + patch_temporal: int = 1, + packing_buffer_size: int = None, + **kwargs, + ): + super().__init__(*args, **kwargs) + self.max_frames = max_frames + self.text_embedding_padding_size = text_embedding_padding_size + self.seq_length = seq_length + self.patch_spatial = patch_spatial + self.patch_temporal = patch_temporal + self.packing_buffer_size = packing_buffer_size + + @abstractmethod + def encode_sample(self, sample: dict) -> dict: + raise NotImplementedError + + def select_samples_to_pack(self, samples: List[DiffusionSample]) -> List[List[DiffusionSample]]: + """ + Selects sequences to pack for mixed image-video training. + """ + results = first_fit_decreasing(samples, self.seq_length) + random.shuffle(results) + return results + + @stateless + def pack_selected_samples(self, samples: List[DiffusionSample]) -> DiffusionSample: + """Construct a new Diffusion sample by concatenating the sequences.""" + + def stack(attr): + return torch.stack([getattr(sample, attr) for sample in samples], dim=0) + + def cat(attr): + return torch.cat([getattr(sample, attr) for sample in samples], dim=0) + + return DiffusionSample( + __key__=",".join([s.__key__ for s in samples]), + __restore_key__=(), # Will be set by energon based on `samples` + __subflavor__=None, + __subflavors__=samples[0].__subflavors__, + video=cat("video"), + context_embeddings=cat("context_embeddings"), + loss_mask=cat("loss_mask"), + seq_len_q=cat("seq_len_q"), + seq_len_kv=cat("seq_len_kv"), + pos_ids=cat("pos_ids"), + latent_shape=stack("latent_shape"), + ) + + @stateless + def batch(self, samples: List[DiffusionSample]) -> dict: + raise NotImplementedError diff --git a/dfm/src/megatron/data/common/sequence_packing_utils.py b/dfm/src/megatron/data/common/sequence_packing_utils.py new file mode 100644 index 00000000..f551a292 --- /dev/null +++ b/dfm/src/megatron/data/common/sequence_packing_utils.py @@ -0,0 +1,105 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. 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 List + + +def find_first_bin_that_fits(bins: List[List[int]], s: int, bin_size: int) -> int: + """ + Finds the first bin in a list of bins that has enough space to fit a sequence of size 's'. + + Args: + bins: A list of lists, where each inner list represents a bin and contains the current elements in that bin. + s: The size of the sequence to be placed in a bin. + bin_size: The maximum capacity of each bin. + + Returns: + The index of the first bin that can fit the sequence 's', or -1 if no such bin exists. + """ + for i, abin in enumerate(bins): + if sum(abin) + s <= bin_size: + return i + return -1 + + +def first_fit(seqlens: List[int], pack_size: int) -> List[List[int]]: + """ + Packs sequences of varying lengths into bins using the First-Fit algorithm. + + Args: + seqlens: A list of integers, representing the lengths of the sequences to be packed. + pack_size: The maximum capacity of each bin. + + Returns: + A list of lists, where each inner list represents a bin and contains the indices + of the sequences assigned to that bin. + """ + res = [] + for s in seqlens: + first_bin = find_first_bin_that_fits(res, s, pack_size) + if first_bin == -1: # open a new bin + res.append([s]) + else: + res[first_bin].append(s) + return res + + +def first_fit_decreasing(seqlens: List[int], pack_size: int) -> List[List[int]]: + """ + Packs sequences of varying lengths into bins using the First-Fit Decreasing algorithm. + + This is a variation of the First-Fit algorithm where the sequences are sorted by decreasing length before packing. + + Args: + seqlens: A list of integers, representing the lengths of the sequences to be packed. + pack_size: The maximum capacity of each bin. + + Returns: + A list of lists, similar to the output of the 'first_fit' function. + """ + sorted_seqlens = sorted(seqlens, reverse=True) + return first_fit(sorted_seqlens, pack_size) + + +def concat_pad(tensor_list, max_seq_length): + """ + Efficiently concatenates a list of tensors along the first dimension and pads with zeros + to reach max_seq_length. + + Args: + tensor_list (list of torch.Tensor): List of tensors to concatenate and pad. + max_seq_length (int): The desired size of the first dimension of the output tensor. + + Returns: + torch.Tensor: A tensor of shape [max_seq_length, ...], where ... represents the remaining dimensions. + """ + import torch + + # Get common properties from the first tensor + other_shape = tensor_list[0].shape[1:] + dtype = tensor_list[0].dtype + device = tensor_list[0].device + + # Initialize the result tensor with zeros + result = torch.zeros((max_seq_length, *other_shape), dtype=dtype, device=device) + + current_index = 0 + for tensor in tensor_list: + length = tensor.shape[0] + # Directly assign the tensor to the result tensor without checks + result[current_index : current_index + length] = tensor + current_index += length + + return result diff --git a/dfm/src/megatron/data/dit/__init__.py b/dfm/src/megatron/data/dit/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/dfm/src/megatron/data/dit/dit_taskencoder.py b/dfm/src/megatron/data/dit/dit_taskencoder.py new file mode 100644 index 00000000..fe3e6180 --- /dev/null +++ b/dfm/src/megatron/data/dit/dit_taskencoder.py @@ -0,0 +1,235 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. 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. + +import warnings +from typing import List + +import torch +import torch.nn.functional as F +from einops import rearrange +from megatron.core import parallel_state +from megatron.energon import SkipSample, stateless + +from dfm.src.megatron.data.common.diffusion_sample import DiffusionSample +from dfm.src.megatron.data.common.diffusion_task_encoder_with_sp import DiffusionTaskEncoderWithSequencePacking + + +class DiTTaskEncoder(DiffusionTaskEncoderWithSequencePacking): + """ + BasicDiffusionTaskEncoder is a class that encodes image/video samples for diffusion tasks. + Attributes: + cookers (list): A list of Cooker objects used for processing. + max_frames (int, optional): The maximum number of frames to consider from the video. Defaults to None. + text_embedding_padding_size (int): The padding size for text embeddings. Defaults to 512. + Methods: + __init__(*args, max_frames=None, text_embedding_padding_size=512, **kwargs): + Initializes the BasicDiffusionTaskEncoder with optional maximum frames and text embedding padding size. + encode_sample(sample: dict) -> dict: + Encodes a given sample dictionary containing video and text data. + Args: + sample (dict): A dictionary containing 'pth' for video latent and 'json' for additional info. + Returns: + dict: A dictionary containing encoded video, text embeddings, text mask, and loss mask. + Raises: + SkipSample: If the video latent contains NaNs, Infs, or is not divisible by the tensor parallel size. + """ + + def __init__( + self, + *args, + **kwargs, + ): + super().__init__(*args, **kwargs) + + @stateless(restore_seeds=True) + def encode_sample(self, sample: dict) -> DiffusionSample: + video_latent = sample["pth"] + + if torch.isnan(video_latent).any() or torch.isinf(video_latent).any(): + raise SkipSample() + if torch.max(torch.abs(video_latent)) > 1e3: + raise SkipSample() + + info = sample["json"] + video_latent = video_latent.squeeze(0) + C, T, H, W = video_latent.shape + seq_len = ( + video_latent.shape[-1] + * video_latent.shape[-2] + * video_latent.shape[-3] + // self.patch_spatial**2 + // self.patch_temporal + ) + is_image = T == 1 + + if seq_len > self.seq_length: + print(f"Skipping sample {sample['__key__']} because seq_len {seq_len} > self.seq_length {self.seq_length}") + raise SkipSample() + + if self.max_frames is not None: + video_latent = video_latent[:, : self.max_frames, :, :] + + tpcp_size = parallel_state.get_tensor_model_parallel_world_size() + if parallel_state.get_context_parallel_world_size() > 1: + tpcp_size *= parallel_state.get_context_parallel_world_size() * 2 + if (T * H * W) % tpcp_size != 0: + warnings.warn(f"skipping {video_latent.shape=} not divisible by {tpcp_size=}") + raise SkipSample() + + video_latent = rearrange( + video_latent, + "C (T pt) (H ph) (W pw) -> (T H W) (ph pw pt C)", + ph=self.patch_spatial, + pw=self.patch_spatial, + pt=self.patch_temporal, + ) + sample["pickle"] = sample["pickle"].cpu().float().numpy().squeeze(0) + if is_image: + t5_text_embeddings = torch.from_numpy(sample["pickle"]).to(torch.bfloat16) + else: + t5_text_embeddings = torch.from_numpy(sample["pickle"][0]).to(torch.bfloat16) + t5_text_embeddings_seq_length = t5_text_embeddings.shape[0] + + if t5_text_embeddings_seq_length > self.text_embedding_padding_size: + t5_text_embeddings = t5_text_embeddings[: self.text_embedding_padding_size] + else: + t5_text_embeddings = F.pad( + t5_text_embeddings, + ( + 0, + 0, + 0, + self.text_embedding_padding_size - t5_text_embeddings_seq_length, + ), + ) + t5_text_mask = torch.ones(t5_text_embeddings_seq_length, dtype=torch.bfloat16) + + if is_image: + h, w = info["image_height"], info["image_width"] + fps = torch.tensor([30] * 1, dtype=torch.bfloat16) + num_frames = torch.tensor([1] * 1, dtype=torch.bfloat16) + else: + h, w = info["height"], info["width"] + fps = torch.tensor([info["framerate"]] * 1, dtype=torch.bfloat16) + num_frames = torch.tensor([info["num_frames"]] * 1, dtype=torch.bfloat16) + image_size = torch.tensor([[h, w, h, w]] * 1, dtype=torch.bfloat16) + + pos_ids = rearrange( + pos_id_3d.get_pos_id_3d(t=T // self.patch_temporal, h=H // self.patch_spatial, w=W // self.patch_spatial), + "T H W d -> (T H W) d", + ) + + if self.packing_buffer_size is None: + pos_ids = F.pad(pos_ids, (0, 0, 0, self.seq_length - seq_len)) + loss_mask = torch.zeros(self.seq_length, dtype=torch.bfloat16) + loss_mask[:seq_len] = 1 + video_latent = F.pad(video_latent, (0, 0, 0, self.seq_length - seq_len)) + else: + loss_mask = torch.ones(seq_len, dtype=torch.bfloat16) + + return DiffusionSample( + __key__=sample["__key__"], + __restore_key__=sample["__restore_key__"], + __subflavor__=None, + __subflavors__=sample["__subflavors__"], + video=video_latent, + context_embeddings=t5_text_embeddings, + context_mask=t5_text_mask, + loss_mask=loss_mask, + seq_len_q=torch.tensor([seq_len], dtype=torch.int32), + seq_len_kv=torch.tensor([self.text_embedding_padding_size], dtype=torch.int32), + pos_ids=pos_ids, + latent_shape=torch.tensor([C, T, H, W], dtype=torch.int32), + ) + + @stateless + def batch(self, samples: List[DiffusionSample]) -> dict: + """Return dictionary with data for batch.""" + if self.packing_buffer_size is None: + # no packing + return super().batch(samples).to_dict() + + # packing + sample = samples[0] + return dict( + video=sample.video.unsqueeze_(0), + context_embeddings=sample.context_embeddings.unsqueeze_(0), + context_mask=sample.context_mask.unsqueeze_(0) if sample.context_mask is not None else None, + loss_mask=sample.loss_mask.unsqueeze_(0) if sample.loss_mask is not None else None, + seq_len_q=sample.seq_len_q, + seq_len_kv=sample.seq_len_kv, + pos_ids=sample.pos_ids.unsqueeze_(0) if sample.pos_ids is not None else None, + latent_shape=sample.latent_shape, + ) + + +class PosID3D: + def __init__(self, *, max_t=32, max_h=128, max_w=128): + self.max_t = max_t + self.max_h = max_h + self.max_w = max_w + self.generate_pos_id() + + def generate_pos_id(self): + self.grid = torch.stack( + torch.meshgrid( + torch.arange(self.max_t, device="cpu"), + torch.arange(self.max_h, device="cpu"), + torch.arange(self.max_w, device="cpu"), + ), + dim=-1, + ) + + def get_pos_id_3d(self, *, t, h, w): + if t > self.max_t or h > self.max_h or w > self.max_w: + self.max_t = max(self.max_t, t) + self.max_h = max(self.max_h, h) + self.max_w = max(self.max_w, w) + self.generate_pos_id() + return self.grid[:t, :h, :w] + + +pos_id_3d = PosID3D() + + +# def cook_raw_iamges(sample: dict) -> dict: +# """ +# Processes a raw sample dictionary from energon dataset and returns a new dictionary with specific keys. + +# Args: +# sample (dict): The input dictionary containing the raw sample data. + +# Returns: +# dict: A new dictionary containing the processed sample data with the following keys: +# - All keys from the result of `basic_sample_keys(sample)` +# - 'jpg': original images +# - 'png': contains control images +# - 'txt': contains raw text +# """ +# return dict( +# **basic_sample_keys(sample), +# images=sample["jpg"], +# hint=sample["png"], +# txt=sample["txt"], +# ) + + +# class RawImageDiffusionTaskEncoder(DefaultTaskEncoder): +# """ +# Dummy task encoder takes raw image input on CrudeDataset. +# """ + +# cookers = [ +# Cooker(cook_raw_iamges), +# ] diff --git a/dfm/src/megatron/model/common/__init__.py b/dfm/src/megatron/model/common/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/dfm/src/megatron/model/common/dit_attention.py b/dfm/src/megatron/model/common/dit_attention.py new file mode 100644 index 00000000..73ffdd76 --- /dev/null +++ b/dfm/src/megatron/model/common/dit_attention.py @@ -0,0 +1,320 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. 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. + +# pylint: disable=C0115,C0116,C0301 + +import copy +from dataclasses import dataclass +from typing import Union + +import torch +from megatron.core import parallel_state, tensor_parallel +from megatron.core.extensions.transformer_engine import SplitAlongDim +from megatron.core.transformer.attention import ( + CrossAttention, + SelfAttention, + SelfAttentionSubmodules, +) +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.transformer_config import TransformerConfig + + +@dataclass +class DiTCrossAttentionSubmodules: + """ + Configuration class for specifying the submodules of a cross-attention. + """ + + linear_q: Union[ModuleSpec, type] = None + linear_kv: Union[ModuleSpec, type] = None + core_attention: Union[ModuleSpec, type] = None + linear_proj: Union[ModuleSpec, type] = None + layernorm_across_head: bool = False + q_layernorm: Union[ModuleSpec, type] = None + k_layernorm: Union[ModuleSpec, type] = None + + +class DiTSelfAttention(SelfAttention): + def __init__( + self, + config: TransformerConfig, + submodules: SelfAttentionSubmodules, + layer_number: int, + attn_mask_type: AttnMaskType, + cp_comm_type: str = None, + pg_collection=None, + ): + super().__init__( + config, + submodules, + layer_number, + attn_mask_type, + cp_comm_type, + pg_collection, + ) + + self.layernorm_across_head = submodules.layernorm_across_head + + # override q_layernorm + if submodules.q_layernorm is not None: + if self.layernorm_across_head: + q_layernorm_size = self.query_projection_size + else: + q_layernorm_size = self.hidden_size_per_attention_head + norm_config = copy.deepcopy(self.config) + norm_config.normalization = "RMSNorm" + self.q_layernorm = build_module( + submodules.q_layernorm, + eps=norm_config.layernorm_epsilon, + hidden_size=q_layernorm_size, + config=norm_config, + ) + else: + self.q_layernorm = None + + # override k_layernorm + if submodules.k_layernorm is not None: + if self.layernorm_across_head: + k_layernorm_size = self.kv_projection_size + else: + k_layernorm_size = self.hidden_size_per_attention_head + norm_config = copy.deepcopy(self.config) + norm_config.normalization = "RMSNorm" + self.k_layernorm = build_module( + submodules.k_layernorm, + eps=norm_config.layernorm_epsilon, + hidden_size=k_layernorm_size, + config=norm_config, + ) + else: + self.k_layernorm = None + + def get_query_key_value_tensors(self, hidden_states, key_value_states=None, split_qkv=False): + """ + Derives `query`, `key` and `value` tensors from `hidden_states`. + """ + # Attention heads [sq, b, h] --> [sq, b, ng * (np/ng + 2) * hn)] + mixed_qkv, _ = self.linear_qkv(hidden_states) + + # [sq, b, hp] --> [sq, b, ng, (np/ng + 2) * hn] + new_tensor_shape = mixed_qkv.size()[:-1] + ( + self.num_query_groups_per_partition, + ( + (self.num_attention_heads_per_partition // self.num_query_groups_per_partition + 2) + * self.hidden_size_per_attention_head + ), + ) + mixed_qkv = mixed_qkv.view(*new_tensor_shape) + + split_arg_list = [ + ( + self.num_attention_heads_per_partition + // self.num_query_groups_per_partition + * self.hidden_size_per_attention_head + ), + self.hidden_size_per_attention_head, + self.hidden_size_per_attention_head, + ] + + if SplitAlongDim is not None: + # [sq, b, ng, (np/ng + 2) * hn] + # --> [sq, b, ng, np/ng * hn], [sq, b, ng, hn], [sq, b, ng, hn] + (query, key, value) = SplitAlongDim(mixed_qkv, 3, split_arg_list) + else: + # [sq, b, ng, (np/ng + 2) * hn] + # --> [sq, b, ng, np/ng * hn], [sq, b, ng, hn], [sq, b, ng, hn] + (query, key, value) = torch.split(mixed_qkv, split_arg_list, dim=3) + + # [sq, b, ng, np/ng * hn] -> [sq, b, np, hn] + query = query.reshape(query.size(0), query.size(1), -1, self.hidden_size_per_attention_head) + + # gather query and key heads across TP ranks if self.layernorm_across_head is True + if self.layernorm_across_head and parallel_state.get_tensor_model_parallel_world_size() > 1: + query = query.transpose(-2, -1) + key = key.transpose(-2, -1) + query = tensor_parallel.gather_from_tensor_model_parallel_region(query) + key = tensor_parallel.gather_from_tensor_model_parallel_region(key) + query = query.transpose(-2, -1) + key = key.transpose(-2, -1) + + if self.q_layernorm is not None: + if self.layernorm_across_head: + q_flat = query.reshape(query.size(0), query.size(1), -1).contiguous() # [sq, b, np*hn] + q_flat = self.q_layernorm(q_flat) + query = q_flat.view( + query.size(0), query.size(1), -1, self.hidden_size_per_attention_head + ) # [sq, b, np, hn] + else: + query = self.q_layernorm(query.contiguous()) + + if self.k_layernorm is not None: + if self.layernorm_across_head: + k_flat = key.reshape(key.size(0), key.size(1), -1).contiguous() + k_flat = self.k_layernorm(k_flat) + key = k_flat.view(key.size(0), key.size(1), -1, self.hidden_size_per_attention_head) + else: + key = self.k_layernorm(key.contiguous()) + + # scatter query and key heads across TP ranks if self.layernorm_across_head is True + if self.layernorm_across_head and parallel_state.get_tensor_model_parallel_world_size() > 1: + query = query.transpose(-2, -1) + key = key.transpose(-2, -1) + query = tensor_parallel.scatter_to_tensor_model_parallel_region(query) + key = tensor_parallel.scatter_to_tensor_model_parallel_region(key) + query = query.transpose(-2, -1) + key = key.transpose(-2, -1) + query = query.contiguous() # important becuase TE attention expects contiguous tensors + key = key.contiguous() # important becuase TE attention expects contiguous tensors + + if self.config.test_mode: + self.run_realtime_tests() + + return query, key, value + + +class DiTCrossAttention(CrossAttention): + def __init__( + self, + config: TransformerConfig, + submodules: DiTCrossAttentionSubmodules, + layer_number: int, + attn_mask_type: AttnMaskType, + cp_comm_type: str = None, + pg_collection=None, + ): + super().__init__( + config, + submodules, + layer_number, + attn_mask_type, + cp_comm_type, + pg_collection, + ) + + self.layernorm_across_head = submodules.layernorm_across_head + + # override q_layernorm + if submodules.q_layernorm is not None: + if self.layernorm_across_head: + q_layernorm_size = self.query_projection_size + else: + q_layernorm_size = self.hidden_size_per_attention_head + norm_config = copy.deepcopy(self.config) + norm_config.normalization = "RMSNorm" + self.q_layernorm = build_module( + submodules.q_layernorm, + eps=norm_config.layernorm_epsilon, + hidden_size=q_layernorm_size, + config=norm_config, + ) + else: + self.q_layernorm = None + + # override k_layernorm + if submodules.k_layernorm is not None: + if self.layernorm_across_head: + k_layernorm_size = self.kv_projection_size + else: + k_layernorm_size = self.hidden_size_per_attention_head + norm_config = copy.deepcopy(self.config) + norm_config.normalization = "RMSNorm" + self.k_layernorm = build_module( + submodules.k_layernorm, + eps=norm_config.layernorm_epsilon, + hidden_size=k_layernorm_size, + config=norm_config, + ) + else: + self.k_layernorm = None + + self.linear_kv = build_module( + submodules.linear_kv, + self.config.crossattn_emb_size, + 2 * self.kv_projection_size, + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=self.config.add_bias_linear, + skip_bias_add=False, + is_expert=False, + ) + + def get_query_key_value_tensors(self, hidden_states, key_value_states, split_qkv=False): + """ + Derives `query` tensor from `hidden_states`, and `key`/`value` tensors + from `key_value_states`. + """ + # Attention heads [sk, b, h] --> [sk, b, (np * 2 * hn)] + mixed_kv, _ = self.linear_kv(key_value_states) + + # [sk, b, (np * 2 * hn)] --> [sk, b, np, 2 * hn] + new_tensor_shape = mixed_kv.size()[:-1] + ( + self.num_attention_heads_per_partition, + 2 * self.hidden_size_per_attention_head, + ) + mixed_kv = mixed_kv.view(*new_tensor_shape) + + # [sk, b, np, 2 * hn] --> 2 [sk, b, np, hn] + (key, value) = tensor_parallel.split_tensor_along_last_dim(mixed_kv, 2) + + # Attention head [sq, b, h] --> [sq, b, hp] + query, _ = self.linear_q(hidden_states) + + # [sq, b, hp] --> [sq, b, np, hn] + new_tensor_shape = query.size()[:-1] + ( + self.num_attention_heads_per_partition, + self.hidden_size_per_attention_head, + ) + query = query.view(*new_tensor_shape) + + # gather query and key heads across TP ranks if self.layernorm_across_head is True + if self.layernorm_across_head and parallel_state.get_tensor_model_parallel_world_size() > 1: + query = query.transpose(-2, -1) + key = key.transpose(-2, -1) + query = tensor_parallel.gather_from_tensor_model_parallel_region(query) + key = tensor_parallel.gather_from_tensor_model_parallel_region(key) + query = query.transpose(-2, -1) + key = key.transpose(-2, -1) + + if self.q_layernorm is not None: + if self.layernorm_across_head: + q_flat = query.reshape(query.size(0), query.size(1), -1).contiguous() # [sq, b, np*hn] + q_flat = self.q_layernorm(q_flat) + query = q_flat.view( + query.size(0), query.size(1), -1, self.hidden_size_per_attention_head + ) # [sq, b, np, hn] + else: + query = self.q_layernorm(query.contiguous()) + + if self.k_layernorm is not None: + if self.layernorm_across_head: + k_flat = key.reshape(key.size(0), key.size(1), -1).contiguous() + k_flat = self.k_layernorm(k_flat) + key = k_flat.view(key.size(0), key.size(1), -1, self.hidden_size_per_attention_head) + else: + key = self.k_layernorm(key.contiguous()) + + # scatter query and key heads across TP ranks if self.layernorm_across_head is True + if self.layernorm_across_head and parallel_state.get_tensor_model_parallel_world_size() > 1: + query = query.transpose(-2, -1) + key = key.transpose(-2, -1) + query = tensor_parallel.scatter_to_tensor_model_parallel_region(query) + key = tensor_parallel.scatter_to_tensor_model_parallel_region(key) + query = query.transpose(-2, -1) + key = key.transpose(-2, -1) + query = query.contiguous() # important becuase TE attention expects contiguous tensors + key = key.contiguous() # important becuase TE attention expects contiguous tensors + + return query, key, value diff --git a/dfm/src/megatron/model/common/dit_embeddings.py b/dfm/src/megatron/model/common/dit_embeddings.py new file mode 100644 index 00000000..496990ca --- /dev/null +++ b/dfm/src/megatron/model/common/dit_embeddings.py @@ -0,0 +1,164 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. 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. + +# pylint: disable=C0115,C0116,C0301 + + +import logging + +import torch +from diffusers.models.embeddings import TimestepEmbedding, get_3d_sincos_pos_embed +from einops import rearrange +from megatron.core import parallel_state +from megatron.core.transformer.module import MegatronModule + + +log = logging.getLogger(__name__) + + +# To be used from Common +class ParallelTimestepEmbedding(TimestepEmbedding): + """ + ParallelTimestepEmbedding is a subclass of TimestepEmbedding that initializes + the embedding layers with an optional random seed for syncronization. + + Args: + in_channels (int): Number of input channels. + time_embed_dim (int): Dimension of the time embedding. + seed (int, optional): Random seed for initializing the embedding layers. + If None, no specific seed is set. + + Attributes: + linear_1 (nn.Module): First linear layer for the embedding. + linear_2 (nn.Module): Second linear layer for the embedding. + + Methods: + __init__(in_channels, time_embed_dim, seed=None): Initializes the embedding layers. + """ + + def __init__(self, in_channels: int, time_embed_dim: int, seed=None): + super().__init__(in_channels=in_channels, time_embed_dim=time_embed_dim) + if seed is not None: + with torch.random.fork_rng(): + torch.manual_seed(seed) + self.linear_1.reset_parameters() + self.linear_2.reset_parameters() + + if parallel_state.get_pipeline_model_parallel_world_size() > 1: + setattr(self.linear_1.weight, "pipeline_parallel", True) + setattr(self.linear_1.bias, "pipeline_parallel", True) + setattr(self.linear_2.weight, "pipeline_parallel", True) + setattr(self.linear_2.bias, "pipeline_parallel", True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Computes the positional embeddings for the input tensor. + + Args: + x (torch.Tensor): Input tensor of shape (B, T, H, W, C). + + Returns: + torch.Tensor: Positional embeddings of shape (B, T, H, W, C). + """ + return super().forward(x.to(torch.bfloat16, non_blocking=True)) + + +class SinCosPosEmb3D(MegatronModule): + """ + SinCosPosEmb3D is a 3D sine-cosine positional embedding module. + + Args: + model_channels (int): Number of channels in the model. + h (int): Length of the height dimension. + w (int): Length of the width dimension. + t (int): Length of the temporal dimension. + spatial_interpolation_scale (float, optional): Scale factor for spatial interpolation. Default is 1.0. + temporal_interpolation_scale (float, optional): Scale factor for temporal interpolation. Default is 1.0. + + Methods: + forward(pos_ids: torch.Tensor) -> torch.Tensor: + Computes the positional embeddings for the input tensor. + + Args: + pos_ids (torch.Tensor): Input tensor of shape (B S 3). + + Returns: + torch.Tensor: Positional embeddings of shape (B S D). + """ + + def __init__( + self, + config, + h: int, + w: int, + t: int, + spatial_interpolation_scale=1.0, + temporal_interpolation_scale=1.0, + ): + super().__init__(config=config) + self.h = h + self.w = w + self.t = t + # h w t + param = get_3d_sincos_pos_embed( + config.hidden_size, [h, w], t, spatial_interpolation_scale, temporal_interpolation_scale, output_type="pt" + ) + param = rearrange(param, "t hw c -> (t hw) c") + self.pos_embedding = torch.nn.Embedding(param.shape[0], config.hidden_size) + self.pos_embedding.weight = torch.nn.Parameter(torch.tensor(param), requires_grad=False) + + def forward(self, pos_ids: torch.Tensor): + # pos_ids: t h w + pos_id = pos_ids[..., 0] * self.h * self.w + pos_ids[..., 1] * self.w + pos_ids[..., 2] + return self.pos_embedding(pos_id) + + +class FactorizedLearnable3DEmbedding(MegatronModule): + def __init__( + self, + config, + t: int, + h: int, + w: int, + **kwargs, + ): + super().__init__(config=config) + self.emb_t = torch.nn.Embedding(t, config.hidden_size) + self.emb_h = torch.nn.Embedding(h, config.hidden_size) + self.emb_w = torch.nn.Embedding(w, config.hidden_size) + + if "seed" in kwargs.keys(): + seed = kwargs["seed"] + with torch.random.fork_rng(): + torch.manual_seed(seed) + if config.perform_initialization: + self.customize_init_param() + else: + self.reset_parameters() + else: + if config.perform_initialization: + self.customize_init_param() + + def customize_init_param(self): + self.config.init_method(self.emb_t.weight) + self.config.init_method(self.emb_h.weight) + self.config.init_method(self.emb_w.weight) + + def reset_parameters(self): + self.emb_t.reset_parameters() + self.emb_h.reset_parameters() + self.emb_w.reset_parameters() + + def forward(self, pos_ids: torch.Tensor): + return self.emb_t(pos_ids[..., 0]) + self.emb_h(pos_ids[..., 1]) + self.emb_w(pos_ids[..., 2]) diff --git a/dfm/src/megatron/model/dit/__init__.py b/dfm/src/megatron/model/dit/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/dfm/src/megatron/model/dit/dit_data_process.py b/dfm/src/megatron/model/dit/dit_data_process.py new file mode 100644 index 00000000..e9e9344c --- /dev/null +++ b/dfm/src/megatron/model/dit/dit_data_process.py @@ -0,0 +1,90 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. 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. + +import torch +from megatron.core.packed_seq_params import PackedSeqParams + + +def dit_data_step(qkv_format, dataloader_iter): + # import pdb;pdb.set_trace() + batch = next(iter(dataloader_iter.iterable)) + batch = get_batch_on_this_cp_rank(batch) + batch = {k: v.to(device="cuda", non_blocking=True) if torch.is_tensor(v) else v for k, v in batch.items()} + batch["is_preprocessed"] = True # assume data is preprocessed + return encode_seq_length(batch, format=qkv_format) + + +def encode_seq_length(batch, format): + if ("seq_len_q" in batch) and ("seq_len_kv" in batch): + zero = torch.zeros([1], dtype=torch.int32, device="cuda") + + cu_seqlens_q = batch["seq_len_q"].cumsum(dim=0).to(torch.int32) + cu_seqlens_q = torch.cat((zero, cu_seqlens_q)) + + cu_seqlens_kv = batch["seq_len_kv"].cumsum(dim=0).to(torch.int32) + cu_seqlens_kv = torch.cat((zero, cu_seqlens_kv)) + + batch["packed_seq_params"] = { + "self_attention": PackedSeqParams( + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_q, + cu_seqlens_q_padded=None, + cu_seqlens_kv_padded=None, + qkv_format=format, + ), + "cross_attention": PackedSeqParams( + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + cu_seqlens_q_padded=None, + cu_seqlens_kv_padded=None, + qkv_format=format, + ), + } + + return batch + + +def get_batch_on_this_cp_rank(data): + """Split the data for context parallelism.""" + from megatron.core import mpu + + cp_size = mpu.get_context_parallel_world_size() + cp_rank = mpu.get_context_parallel_rank() + + t = 16 + if cp_size > 1: + # cp split on seq_length, for video_latent, noise_latent and pos_ids + assert t % cp_size == 0, "t must divisibly by cp_size" + num_valid_tokens_in_ub = None + if "loss_mask" in data and data["loss_mask"] is not None: + num_valid_tokens_in_ub = data["loss_mask"].sum() + + for key, value in data.items(): + if (value is not None) and (key in ["video", "video_latent", "noise_latent", "pos_ids"]): + if len(value.shape) > 5: + value = value.squeeze(0) + B, C, T, H, W = value.shape + if T % cp_size == 0: + # FIXME packed sequencing + data[key] = value.view(B, C, cp_size, T // cp_size, H, W)[:, :, cp_rank, ...].contiguous() + else: + # FIXME packed sequencing + data[key] = value.view(B, C, T, cp_size, H // cp_size, W)[:, :, :, cp_rank, ...].contiguous() + loss_mask = data["loss_mask"] + data["loss_mask"] = loss_mask.view(loss_mask.shape[0], cp_size, loss_mask.shape[1] // cp_size)[ + :, cp_rank, ... + ].contiguous() + data["num_valid_tokens_in_ub"] = num_valid_tokens_in_ub + + return data diff --git a/dfm/src/megatron/model/dit/dit_layer_spec.py b/dfm/src/megatron/model/dit/dit_layer_spec.py new file mode 100644 index 00000000..2b19be2c --- /dev/null +++ b/dfm/src/megatron/model/dit/dit_layer_spec.py @@ -0,0 +1,273 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. 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. + +# pylint: disable=C0115,C0116,C0301 + +import copy +from dataclasses import dataclass +from typing import Literal, Optional, Union + +import torch +import torch.nn as nn +from megatron.core.transformer.attention import ( + SelfAttention, + SelfAttentionSubmodules, +) +from megatron.core.transformer.custom_layers.transformer_engine import ( + TEColumnParallelLinear, + TEDotProductAttention, + TENorm, + TERowParallelLinear, +) +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.identity_op import IdentityOp +from megatron.core.transformer.mlp import MLP, MLPSubmodules +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.transformer_block import TransformerConfig +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules +from megatron.core.utils import make_viewless_tensor + +# to be imported from common +from dfm.src.megatron.model.common.dit_attention import DiTCrossAttention, DiTCrossAttentionSubmodules + + +@dataclass +class DiTWithAdaLNSubmodules(TransformerLayerSubmodules): + full_self_attention: Union[ModuleSpec, type] = IdentityOp + + +class RMSNorm(nn.Module): + def __init__(self, hidden_size: int, config=None, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(hidden_size)) + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + output = self._norm(x.float()).type_as(x) + return output * self.weight + + +class AdaLN(MegatronModule): + """ + Adaptive Layer Normalization Module for DiT. + """ + + def __init__( + self, config: TransformerConfig, n_adaln_chunks=9, use_adaln_lora=True, adaln_lora_dim=256, norm=nn.LayerNorm + ): + super().__init__(config) + if norm == TENorm: + self.ln = norm(config, config.hidden_size, config.layernorm_epsilon) + else: + self.ln = norm(config.hidden_size, elementwise_affine=False, eps=self.config.layernorm_epsilon) + self.n_adaln_chunks = n_adaln_chunks + if use_adaln_lora: + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(config.hidden_size, adaln_lora_dim, bias=False), + nn.Linear(adaln_lora_dim, self.n_adaln_chunks * config.hidden_size, bias=False), + ) + else: + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), nn.Linear(config.hidden_size, self.n_adaln_chunks * config.hidden_size, bias=False) + ) + nn.init.constant_(self.adaLN_modulation[-1].weight, 0) + + setattr(self.adaLN_modulation[-1].weight, "sequence_parallel", config.sequence_parallel) + + def forward(self, timestep_emb): + return self.adaLN_modulation(timestep_emb).chunk(self.n_adaln_chunks, dim=-1) + + def modulate(self, x, shift, scale): + return x * (1 + scale) + shift + + def scale_add(self, residual, x, gate): + return residual + gate * x + + def modulated_layernorm(self, x, shift, scale): + input_layernorm_output = self.ln(x).type_as(x) + return self.modulate(input_layernorm_output, shift, scale) + + def scaled_modulated_layernorm(self, residual, x, gate, shift, scale): + hidden_states = self.scale_add(residual, x, gate) + shifted_pre_mlp_layernorm_output = self.modulated_layernorm(hidden_states, shift, scale) + return hidden_states, shifted_pre_mlp_layernorm_output + + +class DiTLayerWithAdaLN(TransformerLayer): + """A single transformer layer. + + Transformer layer takes input with size [s, b, h] and returns an + output of the same size. + + DiT with Adapative Layer Normalization. + """ + + def __init__( + self, + config: TransformerConfig, + submodules: TransformerLayerSubmodules, + layer_number: int = 1, + hidden_dropout: float = None, + position_embedding_type: Literal["learned_absolute", "rope"] = "learned_absolute", + pg_collection=None, + vp_stage: Optional[int] = None, + ): + def _replace_no_cp_submodules(submodules): + modified_submods = copy.deepcopy(submodules) + modified_submods.cross_attention = IdentityOp + # modified_submods.temporal_self_attention = IdentityOp + return modified_submods + + # Replace any submodules that will have CP disabled and build them manually later after TransformerLayer init. + modified_submods = _replace_no_cp_submodules(submodules) + super().__init__( + config=config, submodules=modified_submods, layer_number=layer_number, hidden_dropout=hidden_dropout + ) + + # Override Cross Attention to disable CP. + # Disable TP Comm overlap as well. Not disabling will attempt re-use of buffer size same as Q and lead to + # incorrect tensor shapes. + if submodules.cross_attention != IdentityOp: + cp_override_config = copy.deepcopy(config) + cp_override_config.context_parallel_size = 1 + cp_override_config.tp_comm_overlap = False + self.cross_attention = build_module( + submodules.cross_attention, + config=cp_override_config, + layer_number=layer_number, + ) + else: + self.cross_attention = None + + self.full_self_attention = build_module( + submodules.full_self_attention, + config=self.config, + layer_number=layer_number, + ) + + self.adaLN = AdaLN(config=self.config, n_adaln_chunks=9 if self.cross_attention else 6) + + def forward( + self, + hidden_states, + attention_mask, + context=None, + context_mask=None, + rotary_pos_emb=None, + rotary_pos_cos=None, + rotary_pos_sin=None, + attention_bias=None, + inference_params=None, + packed_seq_params=None, + sequence_len_offset=None, + inference_context=None, + rotary_pos_cos_sin=None, + ): + timestep_emb = attention_mask + + if self.cross_attention: + shift_full, scale_full, gate_full, shift_ca, scale_ca, gate_ca, shift_mlp, scale_mlp, gate_mlp = ( + self.adaLN(timestep_emb) + ) + else: + shift_full, scale_full, gate_full, shift_mlp, scale_mlp, gate_mlp = self.adaLN(timestep_emb) + + pre_full_attn_layernorm_output_ada = self.adaLN.modulated_layernorm( + hidden_states, shift=shift_full, scale=scale_full + ) + attention_output, _ = self.full_self_attention( + pre_full_attn_layernorm_output_ada, + attention_mask=None, + packed_seq_params=None if packed_seq_params is None else packed_seq_params["self_attention"], + ) + + if self.cross_attention: + hidden_states, pre_cross_attn_layernorm_output_ada = self.adaLN.scaled_modulated_layernorm( + residual=hidden_states, + x=attention_output, + gate=gate_full, + shift=shift_ca, + scale=scale_ca, + ) + attention_output, _ = self.cross_attention( + pre_cross_attn_layernorm_output_ada, + attention_mask=context_mask, + key_value_states=context, + packed_seq_params=None if packed_seq_params is None else packed_seq_params["cross_attention"], + ) + + hidden_states, pre_mlp_layernorm_output_ada = self.adaLN.scaled_modulated_layernorm( + residual=hidden_states, + x=attention_output, + gate=gate_ca if self.cross_attention else gate_full, + shift=shift_mlp, + scale=scale_mlp, + ) + + mlp_output, _ = self.mlp(pre_mlp_layernorm_output_ada) + hidden_states = self.adaLN.scale_add(residual=hidden_states, x=mlp_output, gate=gate_mlp) + + # Jit compiled function creates 'view' tensor. This tensor + # potentially gets saved in the MPU checkpoint function context, + # which rejects view tensors. While making a viewless tensor here + # won't result in memory savings (like the data loader, or + # p2p_communication), it serves to document the origin of this + # 'view' tensor. + output = make_viewless_tensor(inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True) + return output, context + + +def get_dit_adaln_block_with_transformer_engine_spec() -> ModuleSpec: + params = {"attn_mask_type": AttnMaskType.padding} + return ModuleSpec( + module=DiTLayerWithAdaLN, + submodules=DiTWithAdaLNSubmodules( + full_self_attention=ModuleSpec( + module=SelfAttention, + params=params, + submodules=SelfAttentionSubmodules( + linear_qkv=TEColumnParallelLinear, + core_attention=TEDotProductAttention, + linear_proj=TERowParallelLinear, + q_layernorm=RMSNorm, + k_layernorm=RMSNorm, + ), + ), + cross_attention=ModuleSpec( + module=DiTCrossAttention, + params=params, + submodules=DiTCrossAttentionSubmodules( + linear_q=TEColumnParallelLinear, + linear_kv=TEColumnParallelLinear, + core_attention=TEDotProductAttention, + linear_proj=TERowParallelLinear, + q_layernorm=RMSNorm, + k_layernorm=RMSNorm, + ), + ), + mlp=ModuleSpec( + module=MLP, + submodules=MLPSubmodules( + linear_fc1=TEColumnParallelLinear, + linear_fc2=TERowParallelLinear, + ), + ), + ), + ) diff --git a/dfm/src/megatron/model/dit/dit_model.py b/dfm/src/megatron/model/dit/dit_model.py new file mode 100644 index 00000000..e3ae8a29 --- /dev/null +++ b/dfm/src/megatron/model/dit/dit_model.py @@ -0,0 +1,341 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. 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. +# pylint: disable=C0115,C0116,C0301 + +from typing import Dict, Literal, Optional + +import torch +import torch.nn as nn +from diffusers.models.embeddings import Timesteps +from einops import rearrange +from megatron.core import parallel_state, tensor_parallel +from megatron.core.dist_checkpointing.mapping import ShardedStateDict +from megatron.core.models.common.vision_module.vision_module import VisionModule +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.transformer.enums import ModelType +from megatron.core.transformer.transformer_block import TransformerBlock +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.utils import make_sharded_tensor_for_checkpoint +from torch import Tensor + +from dfm.src.megatron.model.common.dit_embeddings import ( + FactorizedLearnable3DEmbedding, + ParallelTimestepEmbedding, + SinCosPosEmb3D, +) +from dfm.src.megatron.model.dit.dit_layer_spec import ( + RMSNorm, +) +from dfm.src.megatron.model.dit.dit_layer_spec import ( + get_dit_adaln_block_with_transformer_engine_spec as DiTLayerWithAdaLNspec, +) + + +class DiTCrossAttentionModel(VisionModule): + """ + DiTCrossAttentionModel is a VisionModule that implements a DiT model with a cross-attention block. + Attributes: + config (TransformerConfig): Configuration for the transformer. + pre_process (bool): Whether to apply pre-processing steps. + post_process (bool): Whether to apply post-processing steps. + fp16_lm_cross_entropy (bool): Whether to use fp16 for cross-entropy loss. + parallel_output (bool): Whether to use parallel output. + position_embedding_type (Literal["learned_absolute", "rope"]): Type of position embedding. + max_img_h (int): Maximum image height. + max_img_w (int): Maximum image width. + max_frames (int): Maximum number of frames. + patch_spatial (int): Spatial patch size. + patch_temporal (int): Temporal patch size. + in_channels (int): Number of input channels. + out_channels (int): Number of output channels. + transformer_decoder_layer_spec (DiTLayerWithAdaLNspec): Specification for the transformer decoder layer. + add_encoder (bool): Whether to add an encoder. + add_decoder (bool): Whether to add a decoder. + share_embeddings_and_output_weights (bool): Whether to share embeddings and output weights. + concat_padding_mask (bool): Whether to concatenate padding mask. + pos_emb_cls (str): Class of position embedding. + model_type (ModelType): Type of the model. + decoder (TransformerBlock): Transformer decoder block. + t_embedder (torch.nn.Sequential): Time embedding layer. + x_embedder (nn.Conv3d): Convolutional layer for input embedding. + pos_embedder (SinCosPosEmb3D): Position embedding layer. + final_layer_linear (torch.nn.Linear): Final linear layer. + affline_norm (RMSNorm): Affine normalization layer. + Methods: + forward(x: Tensor, timesteps: Tensor, crossattn_emb: Tensor, packed_seq_params: PackedSeqParams = None, pos_ids: Tensor = None, **kwargs) -> Tensor: + Forward pass of the model. + set_input_tensor(input_tensor: Tensor) -> None: + Sets input tensor to the model. + sharded_state_dict(prefix: str = 'module.', sharded_offsets: tuple = (), metadata: Optional[Dict] = None) -> ShardedStateDict: + Sharded state dict implementation for backward-compatibility. + tie_embeddings_weights_state_dict(tensor, sharded_state_dict: ShardedStateDict, output_layer_weight_key: str, first_stage_word_emb_key: str) -> None: + Ties the embedding and output weights in a given sharded state dict. + """ + + def __init__( + self, + config: TransformerConfig, + pre_process: bool = True, + post_process: bool = True, + fp16_lm_cross_entropy: bool = False, + parallel_output: bool = True, + position_embedding_type: Literal["learned_absolute", "rope"] = "rope", + max_img_h: int = 80, + max_img_w: int = 80, + max_frames: int = 34, + patch_spatial: int = 1, + patch_temporal: int = 1, + in_channels: int = 16, + out_channels: int = 16, + transformer_decoder_layer_spec=DiTLayerWithAdaLNspec, + pos_embedder=SinCosPosEmb3D, + **kwargs, + ): + super(DiTCrossAttentionModel, self).__init__(config=config) + + self.config: TransformerConfig = config + + self.transformer_decoder_layer_spec = transformer_decoder_layer_spec() + self.pre_process = pre_process + self.post_process = post_process + self.add_encoder = True + self.add_decoder = True + self.fp16_lm_cross_entropy = fp16_lm_cross_entropy + self.parallel_output = parallel_output + self.position_embedding_type = position_embedding_type + self.share_embeddings_and_output_weights = False + self.concat_padding_mask = True + self.pos_emb_cls = "sincos" + self.patch_spatial = patch_spatial + self.patch_temporal = patch_temporal + + # megatron core pipelining currently depends on model type + # TODO: remove this dependency ? + self.model_type = ModelType.encoder_or_decoder + + # Transformer decoder + self.decoder = TransformerBlock( + config=self.config, + spec=self.transformer_decoder_layer_spec, + pre_process=self.pre_process, + post_process=False, + post_layer_norm=False, + ) + + self.t_embedder = torch.nn.Sequential( + Timesteps(self.config.hidden_size, flip_sin_to_cos=False, downscale_freq_shift=0), + ParallelTimestepEmbedding(self.config.hidden_size, self.config.hidden_size, seed=1234), + ) + + self.fps_embedder = nn.Sequential( + Timesteps(num_channels=256, flip_sin_to_cos=False, downscale_freq_shift=1), + ParallelTimestepEmbedding(256, 256, seed=1234), + ) + + if self.pre_process: + self.x_embedder = torch.nn.Linear(in_channels * patch_spatial**2, self.config.hidden_size) + + if pos_embedder is SinCosPosEmb3D: + if self.pre_process: + self.pos_embedder = pos_embedder( + config, + t=max_frames // patch_temporal, + h=max_img_h // patch_spatial, + w=max_img_w // patch_spatial, + ) + else: + # here I just follow the original logic, that except with SinCosPosEmb3D, the pos_emb would be feeded to transformer blocks, + # so the other embedders should be replicated across pp ranks. + self.pos_embedder = pos_embedder( + config, + t=max_frames // patch_temporal, + h=max_img_h // patch_spatial, + w=max_img_w // patch_spatial, + seed=1234, + ) + if parallel_state.get_pipeline_model_parallel_world_size() > 1: + for p in self.pos_embedder.parameters(): + setattr(p, "pipeline_parallel", True) + + if self.post_process: + self.final_layer_linear = torch.nn.Linear( + self.config.hidden_size, + patch_spatial**2 * patch_temporal * out_channels, + ) + + self.affline_norm = RMSNorm(self.config.hidden_size) + if parallel_state.get_pipeline_model_parallel_world_size() > 1: + setattr(self.affline_norm.weight, "pipeline_parallel", True) + + def forward( + self, + x: Tensor, + timesteps: Tensor, + crossattn_emb: Tensor, + packed_seq_params: PackedSeqParams = None, + pos_ids: Tensor = None, + **kwargs, + ) -> Tensor: + """Forward pass. + + Args: + x (Tensor): vae encoded data (b s c) + encoder_decoder_attn_mask (Tensor): cross-attention mask between encoder and decoder + inference_params (InferenceParams): relevant arguments for inferencing + + Returns: + Tensor: loss tensor + """ + B = x.shape[0] + fps = kwargs.get( + "fps", + torch.tensor( + [ + 30, + ] + * B, + dtype=torch.bfloat16, + device=x.device, + ), + ).view(-1) + if self.pre_process: + # transpose to match + x_B_S_D = self.x_embedder(x) + if isinstance(self.pos_embedder, SinCosPosEmb3D): + pos_emb = None + x_B_S_D += self.pos_embedder(pos_ids) + else: + pos_emb = self.pos_embedder(pos_ids) + pos_emb = rearrange(pos_emb, "B S D -> S B D") + x_S_B_D = rearrange(x_B_S_D, "B S D -> S B D") + else: + # intermediate stage of pipeline + x_S_B_D = None ### should it take encoder_hidden_states + if (not hasattr(self, "pos_embedder")) or isinstance(self.pos_embedder, SinCosPosEmb3D): + pos_emb = None + else: + # if transformer blocks need pos_emb, then pos_embedder should + # be replicated across pp ranks. + pos_emb = rearrange(self.pos_embedder(pos_ids), "B S D -> S B D") + + timesteps_B_D = self.t_embedder(timesteps.flatten()).to(torch.bfloat16) # (b d_text_embedding) + + affline_emb_B_D = timesteps_B_D + fps_B_D = self.fps_embedder(fps) + fps_B_D = nn.functional.pad(fps_B_D, (0, self.config.hidden_size - fps_B_D.shape[1])) + affline_emb_B_D += fps_B_D + + crossattn_emb = rearrange(crossattn_emb, "B S D -> S B D") + + if self.config.sequence_parallel: + if self.pre_process: + x_S_B_D = tensor_parallel.scatter_to_sequence_parallel_region(x_S_B_D) + if isinstance(self.pos_embedder, FactorizedLearnable3DEmbedding): + pos_emb = tensor_parallel.scatter_to_sequence_parallel_region(pos_emb) + + crossattn_emb = tensor_parallel.scatter_to_sequence_parallel_region(crossattn_emb) + # `scatter_to_sequence_parallel_region` returns a view, which prevents + # the original tensor from being garbage collected. Clone to facilitate GC. + # Has a small runtime cost (~0.5%). + if self.config.clone_scatter_output_in_embedding: + if self.pre_process: + x_S_B_D = x_S_B_D.clone() + crossattn_emb = crossattn_emb.clone() + + x_S_B_D = self.decoder( + hidden_states=x_S_B_D, + attention_mask=affline_emb_B_D, + context=crossattn_emb, + context_mask=None, + rotary_pos_emb=pos_emb, + packed_seq_params=packed_seq_params, + ) + + if not self.post_process: + return x_S_B_D + + if self.config.sequence_parallel: + x_S_B_D = tensor_parallel.gather_from_sequence_parallel_region(x_S_B_D) + + x_S_B_D = self.final_layer_linear(x_S_B_D) + return rearrange(x_S_B_D, "S B D -> B S D") + + def set_input_tensor(self, input_tensor: Tensor) -> None: + """Sets input tensor to the model. + + See megatron.model.transformer.set_input_tensor() + + Args: + input_tensor (Tensor): Sets the input tensor for the model. + """ + # This is usually handled in schedules.py but some inference code still + # gives us non-lists or None + if not isinstance(input_tensor, list): + input_tensor = [input_tensor] + + assert len(input_tensor) == 1, "input_tensor should only be length 1 for gpt/bert" + self.decoder.set_input_tensor(input_tensor[0]) + + def sharded_state_dict( + self, prefix: str = "module.", sharded_offsets: tuple = (), metadata: Optional[Dict] = None + ) -> ShardedStateDict: + """Sharded state dict implementation for GPTModel backward-compatibility (removing extra state). + + Args: + prefix (str): Module name prefix. + sharded_offsets (tuple): PP related offsets, expected to be empty at this module level. + metadata (Optional[Dict]): metadata controlling sharded state dict creation. + + Returns: + ShardedStateDict: sharded state dict for the GPTModel + """ + sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets, metadata) + + for module in ["t_embedder"]: + for param_name, param in getattr(self, module).named_parameters(): + weight_key = f"{prefix}{module}.{param_name}" + self._set_embedder_weights_replica_id(param, sharded_state_dict, weight_key) + return sharded_state_dict + + def _set_embedder_weights_replica_id( + self, tensor: Tensor, sharded_state_dict: ShardedStateDict, embedder_weight_key: str + ) -> None: + """set replica ids of the weights in t_embedder for sharded state dict. + + Args: + sharded_state_dict (ShardedStateDict): state dict with the weight to tie + weight_key (str): key of the weight in the state dict. + This entry will be replaced with a tied version + + Returns: None, acts in-place + """ + tp_rank = parallel_state.get_tensor_model_parallel_rank() + vpp_rank = parallel_state.get_virtual_pipeline_model_parallel_rank() + vpp_rank = vpp_rank if vpp_rank else 0 + vpp_world = parallel_state.get_virtual_pipeline_model_parallel_world_size() + vpp_world = vpp_world if vpp_world else 1 + pp_rank = parallel_state.get_pipeline_model_parallel_rank() + del sharded_state_dict[embedder_weight_key] + replica_id = ( + tp_rank, + (vpp_rank + pp_rank * vpp_world), + parallel_state.get_data_parallel_rank(with_context_parallel=True), + ) + + sharded_state_dict[embedder_weight_key] = make_sharded_tensor_for_checkpoint( + tensor=tensor, + key=embedder_weight_key, + replica_id=replica_id, + allow_shape_mismatch=False, + ) diff --git a/dfm/src/megatron/model/dit/dit_model_provider.py b/dfm/src/megatron/model/dit/dit_model_provider.py new file mode 100644 index 00000000..36639d62 --- /dev/null +++ b/dfm/src/megatron/model/dit/dit_model_provider.py @@ -0,0 +1,156 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. 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. + +import logging +from dataclasses import dataclass +from typing import Callable + +import torch +from megatron.bridge.models.model_provider import ModelProviderMixin +from megatron.core import parallel_state +from megatron.core.models.common.vision_module.vision_module import VisionModule +from megatron.core.transformer.transformer_config import TransformerConfig + +from dfm.src.common.utils.dynamic_import import dynamic_import +from dfm.src.megatron.model.dit.dit_model import DiTCrossAttentionModel + + +logger = logging.getLogger(__name__) + + +@dataclass +class DiTModelProvider(TransformerConfig, ModelProviderMixin[VisionModule]): + """ + Config for DiT-XL model + """ + + crossattn_emb_size: int = 1024 + add_bias_linear: bool = False + gated_linear_unit: bool = False + + num_layers: int = 28 + hidden_size: int = 1152 + max_img_h: int = 80 + max_img_w: int = 80 + max_frames: int = 34 + patch_spatial: int = 2 + patch_temporal: int = 1 + num_attention_heads: int = 16 + layernorm_epsilon = 1e-6 + normalization = "RMSNorm" + add_bias_linear: bool = False + qk_layernorm_per_head: bool = True + layernorm_zero_centered_gamma: bool = False + + fp16_lm_cross_entropy: bool = False + parallel_output: bool = True + share_embeddings_and_output_weights: bool = True + + # max_position_embeddings: int = 5400 + hidden_dropout: float = 0 + attention_dropout: float = 0 + + bf16: bool = True + params_dtype: torch.dtype = torch.bfloat16 + vae_module: str = None + vae_path: str = None + sigma_data: float = 0.5 + in_channels: int = 16 + layernorm_across_heads: bool = False + + replicated_t_embedder = True + qkv_format: str = "thd" + + seq_length: int = 2048 + vocab_size: int = None + make_vocab_size_divisible_by: int = 128 + + def provide(self, pre_process=None, post_process=None, vp_stage=None) -> DiTCrossAttentionModel: + vp_size = self.virtual_pipeline_model_parallel_size + if vp_size: + p_size = self.pipeline_model_parallel_size + assert (self.num_layers // p_size) % vp_size == 0, ( + "Make sure the number of model chunks is the same across all pipeline stages." + ) + + model = DiTCrossAttentionModel + + return model( + self, + fp16_lm_cross_entropy=self.fp16_lm_cross_entropy, + parallel_output=self.parallel_output, + pre_process=parallel_state.is_pipeline_first_stage(), + post_process=parallel_state.is_pipeline_last_stage(), + max_img_h=self.max_img_h, + max_img_w=self.max_img_w, + max_frames=self.max_frames, + patch_spatial=self.patch_spatial, + ) + + def configure_vae(self): + return dynamic_import(self.vae_module)(self.vae_path) + + +@dataclass +class DiT7BModelProvider(DiTModelProvider): + hidden_size: int = 4096 + max_img_h: int = 240 + max_img_w: int = 240 + max_frames: int = 128 + num_attention_heads: int = 32 + apply_rope_fusion: bool = True # TODO: do we support this? + additional_timestamp_channels = None # TODO: do we support this? + bf16: bool = True + vae_module: str = None + vae_path: str = None + + +@dataclass +class DiT14BModelProvider(DiTModelProvider): + num_layers: int = 36 + hidden_size: int = 5120 + crossattn_emb_size: int = 1024 + max_img_h: int = 240 + max_img_w: int = 240 + max_frames: int = 128 + num_attention_heads: int = 40 + apply_rope_fusion: bool = True + layernorm_zero_centered_gamma: bool = False + additional_timestamp_channels = None + bf16: bool = True + vae_module: str = None + vae_path: str = None + loss_add_logvar: bool = True + + +@dataclass +class DiTLlama30BConfig(DiTModelProvider): + num_layers: int = 48 + hidden_size: int = 6144 + ffn_hidden_size: int = 16384 + num_attention_heads: int = 48 + num_query_groups: int = 8 + gated_linear_unit: int = True + bias_activation_fusion: int = True + activation_func: Callable = torch.nn.functional.silu + layernorm_epsilon: float = 1e-5 + max_frames: int = 128 + max_img_h: int = 240 + max_img_w: int = 240 + init_method_std: float = 0.01 + add_bias_linear: bool = False + seq_length: int = 256 + masked_softmax_fusion: bool = True + persist_layer_norm: bool = True + bias_dropout_fusion: bool = True diff --git a/dfm/src/megatron/model/dit/dit_step.py b/dfm/src/megatron/model/dit/dit_step.py new file mode 100644 index 00000000..9bde05d5 --- /dev/null +++ b/dfm/src/megatron/model/dit/dit_step.py @@ -0,0 +1,166 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. 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. + +import logging +import os +from functools import partial +from typing import Iterable + +import torch +from einops import rearrange +from megatron.bridge.training.losses import masked_next_token_loss +from megatron.bridge.training.state import GlobalState +from megatron.core import parallel_state +from megatron.core.models.gpt import GPTModel +from megatron.core.utils import get_model_config + +from dfm.src.common.tokenizers.cosmos.cosmos1.causal_video_tokenizer import CausalVideoTokenizer +from dfm.src.common.utils.save_video import save_video +from dfm.src.megatron.model.dit.dit_data_process import dit_data_step +from dfm.src.megatron.model.dit.edm.edm_pipeline import EDMPipeline + + +logger = logging.getLogger(__name__) + + +class DITForwardStep: + def __init__(self): + self.diffusion_pipeline = EDMPipeline(sigma_data=0.5) + self.valid = False + self.train = True + self.validation_step = 0 + + def on_validation_start(self, batch, model, step): + C, T, H, W = batch["latent_shape"][0] + latent = self.diffusion_pipeline.generate_samples_from_batch( + model, + batch, + guidance=7, + state_shape=batch["video"].shape, + num_steps=35, + is_negative_prompt=True if "neg_context_embeddings" in batch else False, + ) + latent = latent[0, None, : batch["seq_len_q"][0]] + latent = rearrange( + latent, + "b (T H W) (ph pw pt c) -> b c (T pt) (H ph) (W pw)", + ph=model.config.patch_spatial, + pw=model.config.patch_spatial, + pt=model.config.patch_temporal, + c=C, + T=T // model.config.patch_temporal, + H=H // model.config.patch_spatial, + W=W // model.config.patch_spatial, + ) + + vae = CausalVideoTokenizer.from_pretrained("Cosmos-0.1-Tokenizer-CV4x8x8") + vae.to("cuda") + + decoded_video = (1.0 + vae.decode(latent / self.diffusion_pipeline.sigma_data)).clamp(0, 2) / 2 + decoded_video = (decoded_video * 255).to(torch.uint8).permute(0, 2, 3, 4, 1).cpu().numpy() + rank = torch.distributed.get_rank() + image_folder = "validation_images" + os.makedirs(image_folder, exist_ok=True) + save_video( + grid=decoded_video[0], + fps=decoded_video.shape[1], + H=decoded_video.shape[2], + W=decoded_video.shape[3], + video_save_quality=5, + video_save_path=f"{image_folder}/validation_step={step}_rank={rank}.mp4", + ) + + def __call__( + self, state: GlobalState, data_iterator: Iterable, model: GPTModel, return_schedule_plan: bool = False + ) -> tuple[torch.Tensor, partial]: + """Forward training step. + + Args: + state: Global state for the run + data_iterator: Input data iterator + model: The GPT Model + return_schedule_plan (bool): Whether to return the schedule plan instead of the output tensor + + Returns: + tuple containing the output tensor and the loss function + """ + batch = self.data_process(state, data_iterator, model, return_schedule_plan) + if model.training and self.valid: + self.train = True + self.valid = False + elif (not model.training) and self.train: + self.train = False + self.valid = True + self.validation_step += 1 + self.on_validation_start(batch, model, step=self.validation_step) + return self.forward_step(state, batch, model, return_schedule_plan) + + def data_process( + self, state: GlobalState, data_iterator: Iterable, model: GPTModel, return_schedule_plan: bool = False + ) -> tuple[torch.Tensor, partial]: + timers = state.timers + straggler_timer = state.straggler_timer + + config = get_model_config(model) + + timers("batch-generator", log_level=2).start() + qkv_format = getattr(config, "qkv_format", "sbhd") + with straggler_timer(bdata=True): + batch = dit_data_step(qkv_format, data_iterator) + return batch + + def forward_step(self, state, batch, model, return_schedule_plan: bool = False): + timers = state.timers + timers("batch-generator").stop() + + check_for_nan_in_loss = state.cfg.rerun_state_machine.check_for_nan_in_loss + check_for_spiky_loss = state.cfg.rerun_state_machine.check_for_spiky_loss + # import pdb;pdb.set_trace() + straggler_timer = state.straggler_timer + with straggler_timer: + if parallel_state.is_pipeline_last_stage(): + # Self.diffusion_pipeline should not know anything about the model + # TODO: we need to sepearte the noise ingection process from the pipeline itself + output_batch, loss = self.diffusion_pipeline.training_step(model, batch, 0) + output_tensor = torch.mean(loss, dim=-1) + else: + output_tensor = self.diffusion_pipeline.training_step(model, batch, 0) + + loss = output_tensor + if "loss_mask" not in batch or batch["loss_mask"] is None: + loss_mask = torch.ones_like(loss) + else: + loss_mask = batch["loss_mask"] + loss_function = self._create_loss_function(loss_mask, check_for_nan_in_loss, check_for_spiky_loss) + return output_tensor, loss_function + + def _create_loss_function( + self, loss_mask: torch.Tensor, check_for_nan_in_loss: bool, check_for_spiky_loss: bool + ) -> partial: + """Create a partial loss function with the specified configuration. + + Args: + loss_mask: Used to mask out some portions of the loss + check_for_nan_in_loss: Whether to check for NaN values in the loss + check_for_spiky_loss: Whether to check for spiky loss values + + Returns: + A partial function that can be called with output_tensor to compute the loss + """ + return partial( + masked_next_token_loss, + loss_mask, + check_for_nan_in_loss=check_for_nan_in_loss, + check_for_spiky_loss=check_for_spiky_loss, + ) diff --git a/dfm/src/megatron/model/dit/edm/__init__.py b/dfm/src/megatron/model/dit/edm/__init__.py new file mode 100644 index 00000000..d9155f92 --- /dev/null +++ b/dfm/src/megatron/model/dit/edm/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. 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. diff --git a/dfm/src/megatron/model/dit/edm/edm_pipeline.py b/dfm/src/megatron/model/dit/edm/edm_pipeline.py new file mode 100644 index 00000000..4f2f8ece --- /dev/null +++ b/dfm/src/megatron/model/dit/edm/edm_pipeline.py @@ -0,0 +1,434 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. 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, Callable, Dict, Optional, Tuple + +import numpy as np +import torch +from megatron.core import parallel_state +from torch import Tensor + +from dfm.src.common.utils.batch_ops import batch_mul +from dfm.src.common.utils.torch_split_tensor_for_cp import cat_outputs_cp +from dfm.src.megatron.model.dit.edm.edm_utils import EDMSDE, EDMSampler, EDMScaling + + +class EDMPipeline: + """ + EDMPipeline is a class that implements a diffusion model pipeline for video generation. It includes methods for + initializing the pipeline, encoding and decoding video data, performing training steps, denoising, and generating + samples. + Attributes: + p_mean: Mean for SDE process. + p_std: Standard deviation for SDE process. + sigma_max: Maximum noise level. + sigma_min: Minimum noise level. + _noise_generator: Generator for noise. + _noise_level_generator: Generator for noise levels. + sde: SDE process. + sampler: Sampler for the diffusion model. + scaling: Scaling for EDM. + input_data_key: Key for input video data. + input_image_key: Key for input image data. + tensor_kwargs: Tensor keyword arguments. + loss_reduce: Method for reducing loss. + loss_scale: Scale factor for loss. + aesthetic_finetuning: Aesthetic finetuning parameter. + camera_sample_weight: Camera sample weight parameter. + loss_mask_enabled: Flag for enabling loss mask. + Methods: + noise_level_generator: Returns the noise level generator. + _initialize_generators: Initializes noise and noise-level generators. + encode: Encodes input tensor using the video tokenizer. + decode: Decodes latent tensor using video tokenizer. + training_step: Performs a single training step for the diffusion model. + denoise: Performs denoising on the input noise data, noise level, and condition. + compute_loss_with_epsilon_and_sigma: Computes the loss for training. + get_per_sigma_loss_weights: Returns loss weights per sigma noise level. + get_condition_uncondition: Returns conditioning and unconditioning for classifier-free guidance. + get_x0_fn_from_batch: Creates a function to generate denoised predictions with the sampler. + generate_samples_from_batch: Generates samples based on input data batch. + _normalize_video_databatch_inplace: Normalizes video data in-place on a CUDA device to [-1, 1]. + draw_training_sigma_and_epsilon: Draws training noise (epsilon) and noise levels (sigma). + random_dropout_input: Applies random dropout to the input tensor. + get_data_and_condition: Retrieves data and conditioning for model input. + """ + + def __init__( + self, + vae=None, + p_mean=0.0, + p_std=1.0, + sigma_max=80, + sigma_min=0.0002, + sigma_data=0.5, + seed=1234, + ): + """ + Initializes the EDM pipeline with the given parameters. + + Args: + net: The DiT model. + vae: The Video Tokenizer (optional). + p_mean (float): Mean for the SDE. + p_std (float): Standard deviation for the SDE. + sigma_max (float): Maximum sigma value for the SDE. + sigma_min (float): Minimum sigma value for the SDE. + sigma_data (float): Sigma value for EDM scaling. + seed (int): Random seed for reproducibility. + + Attributes: + vae: The Video Tokenizer. + net: The DiT model. + p_mean (float): Mean for the SDE. + p_std (float): Standard deviation for the SDE. + sigma_max (float): Maximum sigma value for the SDE. + sigma_min (float): Minimum sigma value for the SDE. + sigma_data (float): Sigma value for EDM scaling. + seed (int): Random seed for reproducibility. + _noise_generator: Placeholder for noise generator. + _noise_level_generator: Placeholder for noise level generator. + sde: Instance of EDMSDE initialized with p_mean, p_std, sigma_max, and sigma_min. + sampler: Instance of EDMSampler. + scaling: Instance of EDMScaling initialized with sigma_data. + input_data_key (str): Key for input data. + input_image_key (str): Key for input images. + tensor_kwargs (dict): Tensor keyword arguments for device and dtype. + loss_reduce (str): Method to reduce loss ('mean' or other). + loss_scale (float): Scale factor for loss. + """ + self.vae = vae + + self.p_mean = p_mean + self.p_std = p_std + self.sigma_max = sigma_max + self.sigma_min = sigma_min + self.sigma_data = sigma_data + + self.seed = seed + self._noise_generator = None + self._noise_level_generator = None + + self.sde = EDMSDE(p_mean, p_std, sigma_max, sigma_min) + self.sampler = EDMSampler() + self.scaling = EDMScaling(sigma_data) + + self.input_data_key = "video" + self.input_image_key = "images_1024" + self.tensor_kwargs = {"device": "cuda", "dtype": torch.bfloat16} + self.loss_reduce = "mean" + self.loss_scale = 1.0 + + @property + def noise_level_generator(self): + """ + Generates noise levels for the EDM pipeline. + + Returns: + Callable: A function or generator that produces noise levels. + """ + return self._noise_level_generator + + def _initialize_generators(self): + """ + Initializes the random number generators for noise and noise level. + + This method sets up two generators: + 1. A PyTorch generator for noise, seeded with a combination of the base seed and the data parallel rank. + 2. A NumPy generator for noise levels, seeded similarly but without considering context parallel rank. + + Returns: + None + """ + noise_seed = self.seed + 100 * parallel_state.get_data_parallel_rank(with_context_parallel=True) + noise_level_seed = self.seed + 100 * parallel_state.get_data_parallel_rank(with_context_parallel=False) + self._noise_generator = torch.Generator(device="cuda") + self._noise_generator.manual_seed(noise_seed) + self._noise_level_generator = np.random.default_rng(noise_level_seed) + self.sde._generator = self._noise_level_generator + + def training_step( + self, model, data_batch: dict[str, torch.Tensor], iteration: int + ) -> tuple[dict[str, torch.Tensor], torch.Tensor]: + """ + Performs a single training step for the diffusion model. + + This method is responsible for executing one iteration of the model's training. It involves: + 1. Adding noise to the input data using the SDE process. + 2. Passing the noisy data through the network to generate predictions. + 3. Computing the loss based on the difference between the predictions and the original data. + + Args: + data_batch (dict): raw data batch draw from the training data loader. + iteration (int): Current iteration number. + + Returns: + A tuple with the output batch and the computed loss. + """ + # import pdb; pdb.set_trace() + # Get the input data to noise and denoise~(image, video) and the corresponding conditioner. + self.net = model + x0_from_data_batch, x0, condition = self.get_data_and_condition(data_batch) + + # Sample pertubation noise levels and N(0, 1) noises + sigma, epsilon = self.draw_training_sigma_and_epsilon(x0.size(), condition) + + if parallel_state.is_pipeline_last_stage(): + output_batch, pred_mse, edm_loss = self.compute_loss_with_epsilon_and_sigma( + data_batch, x0_from_data_batch, x0, condition, epsilon, sigma + ) + + return output_batch, edm_loss + else: + net_output = self.compute_loss_with_epsilon_and_sigma( + data_batch, x0_from_data_batch, x0, condition, epsilon, sigma + ) + return net_output + + def denoise(self, xt: torch.Tensor, sigma: torch.Tensor, condition: dict[str, torch.Tensor]): + """ + Performs denoising on the input noise data, noise level, and condition + + Args: + xt (torch.Tensor): The input noise data. + sigma (torch.Tensor): The noise level. + condition (dict[str, torch.Tensor]): conditional information + + Returns: + Predicted clean data (x0) and noise (eps_pred). + """ + + xt = xt.to(**self.tensor_kwargs) + sigma = sigma.to(**self.tensor_kwargs) + # get precondition for the network + c_skip, c_out, c_in, c_noise = self.scaling(sigma=sigma) + + net_output = self.net( + x=batch_mul(c_in, xt), # Eq. 7 of https://arxiv.org/pdf/2206.00364.pdf + timesteps=c_noise, # Eq. 7 of https://arxiv.org/pdf/2206.00364.pdf + **condition, + ) + + if not parallel_state.is_pipeline_last_stage(): + return net_output + + x0_pred = batch_mul(c_skip, xt) + batch_mul(c_out, net_output) + + # get noise prediction based on sde + eps_pred = batch_mul(xt - x0_pred, 1.0 / sigma) + + return x0_pred, eps_pred + + def compute_loss_with_epsilon_and_sigma( + self, + data_batch: dict[str, torch.Tensor], + x0_from_data_batch: torch.Tensor, + x0: torch.Tensor, + condition: dict[str, torch.Tensor], + epsilon: torch.Tensor, + sigma: torch.Tensor, + ): + """ + Computes the loss for training. + + Args: + data_batch: Batch of input data. + x0_from_data_batch: Raw input tensor. + x0: Latent tensor. + condition: Conditional input data. + epsilon: Noise tensor. + sigma: Noise level tensor. + + Returns: + The computed loss. + """ + # Get the mean and stand deviation of the marginal probability distribution. + mean, std = self.sde.marginal_prob(x0, sigma) + # Generate noisy observations + xt = mean + batch_mul(std, epsilon) # corrupted data + + if parallel_state.is_pipeline_last_stage(): + # make prediction + x0_pred, eps_pred = self.denoise(xt, sigma, condition) + # loss weights for different noise levels + weights_per_sigma = self.get_per_sigma_loss_weights(sigma=sigma) + pred_mse = (x0 - x0_pred) ** 2 + edm_loss = batch_mul(pred_mse, weights_per_sigma) + + output_batch = { + "x0": x0, + "xt": xt, + "sigma": sigma, + "weights_per_sigma": weights_per_sigma, + "condition": condition, + "model_pred": {"x0_pred": x0_pred, "eps_pred": eps_pred}, + "mse_loss": pred_mse.mean(), + "edm_loss": edm_loss.mean(), + } + return output_batch, pred_mse, edm_loss + else: + # make prediction + x0_pred = self.denoise(xt, sigma, condition) + return x0_pred.contiguous() + + def get_per_sigma_loss_weights(self, sigma: torch.Tensor): + """ + Args: + sigma (tensor): noise level + + Returns: + loss weights per sigma noise level + """ + return (sigma**2 + self.sigma_data**2) / (sigma * self.sigma_data) ** 2 + + def get_condition_uncondition(self, data_batch: Dict): + """Returns conditioning and unconditioning for classifier-free guidance.""" + _, _, condition = self.get_data_and_condition(data_batch, dropout_rate=0.0) + + if "neg_context_embeddings" in data_batch: + data_batch["context_embeddings"] = data_batch["neg_context_embeddings"] + data_batch["context_mask"] = data_batch["context_mask"] + _, _, uncondition = self.get_data_and_condition(data_batch, dropout_rate=1.0) + else: + _, _, uncondition = self.get_data_and_condition(data_batch, dropout_rate=1.0) + + return condition, uncondition + + def get_x0_fn_from_batch( + self, + data_batch: Dict, + guidance: float = 1.5, + is_negative_prompt: bool = False, + ) -> Callable: + """ + Creates a function to generate denoised predictions with the sampler. + + Args: + data_batch: Batch of input data. + guidance: Guidance scale factor. + is_negative_prompt: Whether to use negative prompts. + + Returns: + A callable to predict clean data (x0). + """ + condition, uncondition = self.get_condition_uncondition(data_batch) + + def x0_fn(noise_x: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor: + cond_x0, _ = self.denoise(noise_x, sigma, condition) + uncond_x0, _ = self.denoise(noise_x, sigma, uncondition) + return cond_x0 + guidance * (cond_x0 - uncond_x0) + + return x0_fn + + def generate_samples_from_batch( + self, + model, + data_batch: Dict, + guidance: float = 1.5, + state_shape: Tuple | None = None, + is_negative_prompt: bool = False, + num_steps: int = 35, + ) -> Tensor: + """ + Generates samples based on input data batch. + + Args: + data_batch: Batch of input data. + guidance: Guidance scale factor. + state_shape: Shape of the state. + is_negative_prompt: Whether to use negative prompts. + num_steps: Number of steps for sampling. + solver_option: SDE Solver option. + + Returns: + Generated samples from diffusion model. + """ + self.net = model + cp_enabled = parallel_state.get_context_parallel_world_size() > 1 + + if self._noise_generator is None: + self._initialize_generators() + x0_fn = self.get_x0_fn_from_batch(data_batch, guidance, is_negative_prompt=is_negative_prompt) + + state_shape = list(state_shape) + state_shape[1] //= parallel_state.get_context_parallel_world_size() + x_sigma_max = ( + torch.randn(state_shape, **self.tensor_kwargs, generator=self._noise_generator) * self.sde.sigma_max + ) + + samples = self.sampler(x0_fn, x_sigma_max, num_steps=num_steps, sigma_max=self.sde.sigma_max) + + if cp_enabled: + cp_group = parallel_state.get_context_parallel_group() + samples = cat_outputs_cp(samples, seq_dim=2, cp_group=cp_group) + + return samples + + def draw_training_sigma_and_epsilon(self, x0_size: int, condition: Any) -> torch.Tensor: + """ + Draws training noise (epsilon) and noise levels (sigma). + + Args: + x0_size: Shape of the input tensor. + condition: Conditional input (unused). + + Returns: + Noise level (sigma) and noise (epsilon). + """ + del condition + batch_size = x0_size[0] + if self._noise_generator is None: + self._initialize_generators() + epsilon = torch.randn(x0_size, **self.tensor_kwargs, generator=self._noise_generator) + return self.sde.sample_t(batch_size).to(**self.tensor_kwargs), epsilon + + def random_dropout_input(self, in_tensor: torch.Tensor, dropout_rate: Optional[float] = None) -> torch.Tensor: + """ + Applies random dropout to the input tensor. + + Args: + in_tensor: Input tensor. + dropout_rate: Dropout probability (optional). + + Returns: + Conditioning with random dropout applied. + """ + dropout_rate = dropout_rate if dropout_rate is not None else self.dropout_rate + return batch_mul( + torch.bernoulli((1.0 - dropout_rate) * torch.ones(in_tensor.shape[0])).type_as(in_tensor), + in_tensor, + ) + + def get_data_and_condition(self, data_batch: dict[str, Tensor], dropout_rate=0.2) -> Tuple[Tensor]: + """ + Retrieves data and conditioning for model input. + + Args: + data_batch: Batch of input data. + dropout_rate: Dropout probability for conditioning. + + Returns: + Raw data, latent data, and conditioning information. + """ + # Latent state + raw_state = data_batch["video"] * self.sigma_data + # assume data is already encoded + latent_state = raw_state + + # Condition + data_batch["crossattn_emb"] = self.random_dropout_input( + data_batch["context_embeddings"], dropout_rate=dropout_rate + ) + + return raw_state, latent_state, data_batch diff --git a/dfm/src/megatron/model/dit/edm/edm_utils.py b/dfm/src/megatron/model/dit/edm/edm_utils.py new file mode 100644 index 00000000..698acbb1 --- /dev/null +++ b/dfm/src/megatron/model/dit/edm/edm_utils.py @@ -0,0 +1,137 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. 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. + +# pylint: disable=C0115,C0116,C0301 + +from statistics import NormalDist +from typing import Callable, Tuple + +import numpy as np +import torch +from torch import nn +from tqdm import tqdm + + +class EDMScaling: + def __init__(self, sigma_data: float = 0.5): + self.sigma_data = sigma_data + + def __call__(self, sigma: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + c_skip = self.sigma_data**2 / (sigma**2 + self.sigma_data**2) + c_out = sigma * self.sigma_data / (sigma**2 + self.sigma_data**2) ** 0.5 + c_in = 1 / (sigma**2 + self.sigma_data**2) ** 0.5 + c_noise = 0.25 * sigma.log() + return c_skip, c_out, c_in, c_noise + + +class EDMSDE: + def __init__( + self, + p_mean: float = -1.2, + p_std: float = 1.2, + sigma_max: float = 80.0, + sigma_min: float = 0.002, + ): + self.gaussian_dist = NormalDist(mu=p_mean, sigma=p_std) + self.sigma_max = sigma_max + self.sigma_min = sigma_min + self._generator = np.random + + def sample_t(self, batch_size: int) -> torch.Tensor: + cdf_vals = self._generator.uniform(size=(batch_size)) + samples_interval_gaussian = [self.gaussian_dist.inv_cdf(cdf_val) for cdf_val in cdf_vals] + log_sigma = torch.tensor(samples_interval_gaussian, device="cuda") + return torch.exp(log_sigma) + + def marginal_prob(self, x0: torch.Tensor, sigma: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return x0, sigma + + +class EDMSampler(nn.Module): + """ + Elucidating the Design Space of Diffusion-Based Generative Models (EDM) + # https://github.com/NVlabs/edm/blob/62072d2612c7da05165d6233d13d17d71f213fee/generate.py#L25 + + Attributes: + None + + Methods: + forward(x0_fn: Callable, x_sigma_max: torch.Tensor, num_steps: int = 35, sigma_min: float = 0.002, + sigma_max: float = 80, rho: float = 7, S_churn: float = 0, S_min: float = 0, + S_max: float = float("inf"), S_noise: float = 1) -> torch.Tensor: + Performs the forward pass for the EDM sampling process. + + Parameters: + x0_fn (Callable): A function that takes in a tensor and returns a denoised tensor. + x_sigma_max (torch.Tensor): The initial noise level tensor. + num_steps (int, optional): The number of sampling steps. Default is 35. + sigma_min (float, optional): The minimum noise level. Default is 0.002. + sigma_max (float, optional): The maximum noise level. Default is 80. + rho (float, optional): The rho parameter for time step discretization. Default is 7. + S_churn (float, optional): The churn parameter for noise increase. Default is 0. + S_min (float, optional): The minimum value for the churn parameter. Default is 0. + S_max (float, optional): The maximum value for the churn parameter. Default is float("inf"). + S_noise (float, optional): The noise scale for the churn parameter. Default is 1. + + Returns: + torch.Tensor: The sampled tensor after the EDM process. + """ + + @torch.no_grad() + def forward( + self, + x0_fn: Callable, + x_sigma_max: torch.Tensor, + num_steps: int = 35, + sigma_min: float = 0.002, + sigma_max: float = 80, + rho: float = 7, + S_churn: float = 0, + S_min: float = 0, + S_max: float = float("inf"), + S_noise: float = 1, + ) -> torch.Tensor: + # Time step discretization. + in_dtype = x_sigma_max.dtype + _ones = torch.ones(x_sigma_max.shape[0], dtype=in_dtype, device=x_sigma_max.device) + step_indices = torch.arange(num_steps, dtype=torch.float64, device=x_sigma_max.device) + t_steps = ( + sigma_max ** (1 / rho) + step_indices / (num_steps - 1) * (sigma_min ** (1 / rho) - sigma_max ** (1 / rho)) + ) ** rho + t_steps = torch.cat([t_steps, torch.zeros_like(t_steps[:1])]) # t_N = 0 + + # Main sampling loop. + x_next = x_sigma_max.to(torch.float64) + for i, (t_cur, t_next) in enumerate( + tqdm(zip(t_steps[:-1], t_steps[1:], strict=False), total=len(t_steps) - 1) + ): # 0, ..., N-1 + x_cur = x_next + + # Increase noise temporarily. + gamma = min(S_churn / num_steps, np.sqrt(2) - 1) if S_min <= t_cur <= S_max else 0 + t_hat = t_cur + gamma * t_cur + x_hat = x_cur + (t_hat**2 - t_cur**2).sqrt() * S_noise * torch.randn_like(x_cur) + + # Euler step. + denoised = x0_fn(x_hat.to(in_dtype), t_hat.to(in_dtype) * _ones).to(torch.float64) + d_cur = (x_hat - denoised) / t_hat + x_next = x_hat + (t_next - t_hat) * d_cur + + # Apply 2nd order correction. + if i < num_steps - 1: + denoised = x0_fn(x_hat.to(in_dtype), t_hat.to(in_dtype) * _ones).to(torch.float64) + d_prime = (x_next - denoised) / t_next + x_next = x_hat + (t_next - t_hat) * (0.5 * d_cur + 0.5 * d_prime) + + return x_next.to(in_dtype) diff --git a/dfm/src/megatron/recipes/__init__.py b/dfm/src/megatron/recipes/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/dfm/src/megatron/recipes/dit/__init__.py b/dfm/src/megatron/recipes/dit/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/dfm/src/megatron/recipes/dit/dit.py b/dfm/src/megatron/recipes/dit/dit.py new file mode 100644 index 00000000..14c15e89 --- /dev/null +++ b/dfm/src/megatron/recipes/dit/dit.py @@ -0,0 +1,212 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. 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. + +import os +from typing import List, Optional, Union + +import torch +from megatron.bridge.recipes.utils.optimizer_utils import distributed_fused_adam_with_cosine_annealing +from megatron.bridge.recipes.utils.tokenizer_utils import DEFAULT_NULL_TOKENIZER_VOCAB_SIZE +from megatron.bridge.training.comm_overlap import CommOverlapConfig +from megatron.bridge.training.config import ( + CheckpointConfig, + ConfigContainer, + LoggerConfig, + RNGConfig, + TokenizerConfig, + TrainingConfig, +) +from megatron.bridge.training.mixed_precision import MixedPrecisionConfig, get_mixed_precision_config +from megatron.core.distributed import DistributedDataParallelConfig + +from dfm.src.megatron.data.common.diffusion_energon_datamodule import DiffusionDataModuleConfig +from dfm.src.megatron.model.dit.dit_model_provider import DiTModelProvider + + +def model_config( + tensor_parallelism: int = 1, + pipeline_parallelism: int = 1, + pipeline_parallelism_dtype: Optional[torch.dtype] = torch.bfloat16, + virtual_pipeline_parallelism: Optional[int] = 1, + context_parallelism: int = 1, + sequence_parallelism: bool = False, +) -> DiTModelProvider: + """ + Configure the DiT-S model. + + Args: + tensor_parallelism (int): Degree of tensor model parallelism. + pipeline_parallelism (int): Degree of pipeline model parallelism. + pipeline_parallelism_dtype (Optional[torch.dtype]): Data type for pipeline parallelism. + virtual_pipeline_parallelism (Optional[int]): Size of virtual pipeline parallelism. + context_parallelism (int): Degree of context parallelism. + sequence_parallelism (bool): Whether to use sequence parallelism. + + Returns: + DiTModelProvider: Configuration for the DiT-S model. + """ + return DiTModelProvider( + tensor_model_parallel_size=tensor_parallelism, + pipeline_model_parallel_size=pipeline_parallelism, + pipeline_dtype=pipeline_parallelism_dtype, + virtual_pipeline_model_parallel_size=None, + context_parallel_size=context_parallelism, + sequence_parallel=sequence_parallelism, + seq_length=2048, + ) + + +def pretrain_config( + dir: Optional[str] = None, + name: str = "default", + # Dataset configuration + dataset_path: str = None, + data_paths: Optional[List[str]] = None, + data_args_path: Optional[str] = None, + train_data_path: Optional[List[str]] = None, + valid_data_path: Optional[List[str]] = None, + test_data_path: Optional[List[str]] = None, + per_split_data_args_path: Optional[str] = None, + mock: bool = False, + # Model configuration + tensor_parallelism: int = 1, + pipeline_parallelism: int = 1, + pipeline_parallelism_dtype: Optional[torch.dtype] = torch.bfloat16, + virtual_pipeline_parallelism: Optional[int] = 1, + context_parallelism: int = 1, + sequence_parallelism: bool = False, + use_megatron_fsdp: bool = False, + # Training hyperparameters + train_iters: int = 10000, + global_batch_size: int = 2, # TODO: set it to num devices + micro_batch_size: int = 1, + lr: float = 0.9e-4, + lr_warmup_iters: int = 2000, + # Precision recipe + precision_config: Optional[Union[MixedPrecisionConfig, str]] = "bf16_mixed", + comm_overlap_config: Optional[CommOverlapConfig] = None, +) -> ConfigContainer: + """ + Create a pre-training configuration for GPT3 175B model. + + The default configuration is expected to run on 64 nodes with 8 GPUs each. + + Args: + dir (Optional[str]): Base directory for saving logs and checkpoints. + name (str): Name of the pre-training run. + dataset_path (str): Path to the dataset directory for DiffusionDataModuleConfig. + data_paths (Optional[List[str]]): List of paths to dataset files. If None, mock data will be used. + data_args_path (Optional[str]): Path to file containing data arguments. + train_data_path (Optional[List[str]]): List of training data paths. + valid_data_path (Optional[List[str]]): List of validation data paths. + test_data_path (Optional[List[str]]): List of test data paths. + per_split_data_args_path (Optional[str]): Path to JSON file with per-split data configuration. + mock (bool): Whether to use mock data. If True, ignores data_paths. + tensor_parallelism (int): Degree of tensor model parallelism. + pipeline_parallelism (int): Degree of pipeline model parallelism. + pipeline_parallelism_dtype (Optional[torch.dtype]): Data type for pipeline parallelism. + virtual_pipeline_parallelism (Optional[int]): Size of virtual pipeline parallelism. + context_parallelism (int): Degree of context parallelism to be passed to model_config. + sequence_parallelism (bool): Whether to use sequence parallelism. + train_iters (int): Total number of training iterations. + global_batch_size (int): Global batch size for training. + micro_batch_size (int): Micro batch size for training. + seq_length (int): Sequence length for training data. + lr (float): Learning rate. + min_lr (float): Minimum learning rate for cosine decay. + lr_warmup_iters (int): Number of warmup iterations for the learning rate. + precision_config (Optional[Union[MixedPrecisionConfig, str]]): Precision configuration for the model. + comm_overlap_config (Optional[CommOverlapConfig]): Communication overlap configuration for the model. + + Returns: + ConfigContainer: Configuration for pre-training. + """ + base_output_dir = dir if dir is not None else os.path.join(os.getcwd(), "nemo_experiments") + run_output_dir = os.path.join(base_output_dir, name) + checkpoint_dir = os.path.join(run_output_dir, "checkpoints") + tensorboard_dir = os.path.join(run_output_dir, "tb_logs") + + model_cfg = model_config( + tensor_parallelism=tensor_parallelism, + pipeline_parallelism=pipeline_parallelism, + pipeline_parallelism_dtype=pipeline_parallelism_dtype, + virtual_pipeline_parallelism=virtual_pipeline_parallelism, + context_parallelism=context_parallelism, + sequence_parallelism=sequence_parallelism, + ) + + opt_config, scheduler = distributed_fused_adam_with_cosine_annealing( + lr_warmup_iters=lr_warmup_iters, + lr_decay_iters=train_iters, + max_lr=lr, + ) + opt_config.use_precision_aware_optimizer = False + + if isinstance(precision_config, str): + precision_config = get_mixed_precision_config(precision_config) + + precision_config.grad_reduce_in_fp32 = False + + # Config Container + cfg = ConfigContainer( + model=model_cfg, + train=TrainingConfig( + train_iters=train_iters, + eval_interval=1000, + eval_iters=32, + global_batch_size=global_batch_size, + micro_batch_size=micro_batch_size, + manual_gc=True, + manual_gc_interval=100, + manual_gc_eval=100, + ), + optimizer=opt_config, + scheduler=scheduler, + ddp=DistributedDataParallelConfig( + check_for_nan_in_grad=True, + grad_reduce_in_fp32=True, + overlap_grad_reduce=False, + overlap_param_gather=False, + average_in_collective=True, + use_distributed_optimizer=True, + use_megatron_fsdp=use_megatron_fsdp, # need use_distributed_optimizer=True + ), + dataset=DiffusionDataModuleConfig( + path=dataset_path, + seq_length=2048, + task_encoder_seq_length=8000, + packing_buffer_size=40, + micro_batch_size=micro_batch_size, + global_batch_size=global_batch_size, + num_workers=10, + ), + logger=LoggerConfig( + log_interval=10, + tensorboard_dir=tensorboard_dir, + log_timers_to_tensorboard=True, + ), + tokenizer=TokenizerConfig(tokenizer_type="NullTokenizer", vocab_size=DEFAULT_NULL_TOKENIZER_VOCAB_SIZE), + checkpoint=CheckpointConfig( + save_interval=2000, + save=checkpoint_dir, + load=checkpoint_dir, + ckpt_format="torch_dist", + fully_parallel_save=True, + ), + rng=RNGConfig(seed=1234), + comm_overlap=comm_overlap_config, + mixed_precision=precision_config, + ) + + return cfg diff --git a/examples/megatron/recipes/README.md b/examples/megatron/recipes/README.md new file mode 100644 index 00000000..7ed591b1 --- /dev/null +++ b/examples/megatron/recipes/README.md @@ -0,0 +1,3 @@ +# Recipe + +Training recipes for Wan2.1 pretraining, finetuning, and weight verification. diff --git a/examples/megatron/recipes/__init__.py b/examples/megatron/recipes/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/megatron/recipes/dit/README.md b/examples/megatron/recipes/dit/README.md new file mode 100644 index 00000000..baef566f --- /dev/null +++ b/examples/megatron/recipes/dit/README.md @@ -0,0 +1,77 @@ +# DiT (Diffusion Transformer) Model Setup + +This guide provides instructions for setting up and running the DiT model on the butterfly dataset. + +## Overview + +Megatron-LM and Megatron-Bridge coming with the docker image are not compatible with DiT model. This setup guide will walk you through configuring the environment properly. + +## Setup Instructions + +### 1. Clone Required Repositories + +The following repositories need to be cloned with specific commit hashes: + +#### Megatron-LM +```bash +git clone https://github.com/NVIDIA/Megatron-LM.git +cd Megatron-LM +git checkout aecce9e95624ddfedbd2bd3ce599e36cd96da065 +cd .. +``` + +#### Megatron-Bridge +```bash +git clone https://github.com/NVIDIA-NeMo/Megatron-Bridge.git +cd Megatron-Bridge +git checkout 83f90524f9bc467e8f864e2f9bd1da1246594ab9 +cd .. +``` + +#### DFM Repository +```bash +git clone https://github.com/NVIDIA-NeMo/DFM.git +cd DFM +git checkout dit_debug +cd .. +``` + +### 2. Dataset Location + +The butterfly webdataset is accesible on eos clusters in the path below: +``` +/home/snorouzi/code/butterfly_webdataset +``` + +## Docker Setup + +Run the following Docker command to start the container with all necessary volume mounts: + +```bash +sudo docker run --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 --gpus all -w /opt/dfm --rm \ + -v ${DATA_PATH}/butterfly_webdataset:/opt/VFM/butterfly_webdataset \ + -v ${CODE_PATH}/Megatron-LM:/opt/megatron-lm \ + -v ${CODE_PATH}/Megatron-Bridge/:/opt/Megatron-Bridge/ \ + -v ${CODE_PATH}/DFM:/opt/dfm \ + -it nvcr.io/nvidian/nemo:25.09.rc6 +``` + +**Note:** Set the `DATA_PATH` and `CODE_PATH` environment variables to point to your local directories before running this command. + +## Installation + +Once inside the container, install the required Python packages: + +```bash +pip install --upgrade transformers +pip install imageio==2.24 +pip install imageio[ffmpeg] +``` + +## Running the Model + +Execute the DiT model training with the following command: + +```bash +torchrun --nproc-per-node 2 examples/megatron/recipes/dit/pretrain_dit_model.py --dataset_path "/opt/VFM/butterfly_webdataset" +``` diff --git a/examples/megatron/recipes/dit/__init__.py b/examples/megatron/recipes/dit/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/megatron/recipes/dit/inference_dit_model.py b/examples/megatron/recipes/dit/inference_dit_model.py new file mode 100644 index 00000000..edcf50a6 --- /dev/null +++ b/examples/megatron/recipes/dit/inference_dit_model.py @@ -0,0 +1,299 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +import argparse + +import numpy as np +import torch +from einops import rearrange +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from nemo.lightning.megatron_parallel import MegatronParallel +from nemo_vfm.diffusion.utils.mcore_parallel_utils import Utils +from transformers import T5EncoderModel, T5TokenizerFast + +from dfm.src.common.tokenizers.cosmos.cosmos1.causal_video_tokenizer import CausalVideoTokenizer +from dfm.src.common.utils.save_video import save_video +from dfm.src.megatron.model.dit.dit_model_provider import DiTModelProvider +from dfm.src.megatron.model.dit.edm.edm_pipeline import EDMPipeline + + +MegatronParallel.init_ddp = lambda self: None + +EXAMPLE_PROMPT = ( + "The teal robot is cooking food in a kitchen. Steam rises from a simmering pot " + "as the robot chops vegetables on a worn wooden cutting board. Copper pans hang " + "from an overhead rack, catching glints of afternoon light, while a well-loved " + "cast iron skillet sits on the stovetop next to scattered measuring spoons and " + "a half-empty bottle of olive oil." +) + + +def parse_args(): + parser = argparse.ArgumentParser(description="Video foundation model inference") + parser.add_argument( + "--prompt", + type=str, + default=EXAMPLE_PROMPT, + help="Prompt which the sampled video condition on", + ) + # We turn on negative prompt by default. set to "" to turn it off. + parser.add_argument( + "--negative_prompt", + type=str, + default=None, + help="Negative prompt which the sampled video condition on", + ) + parser.add_argument("--subject_name", type=str, default="", help="Name of fine-tuned subject") + parser.add_argument("--guidance", type=float, default=7, help="Classifier-free guidance scale") + parser.add_argument("--sampler", type=str, default="RES", help="Currently only supports RES sampler.") + parser.add_argument("--video_save_path", type=str, default="outputs", help="Path to save the video") + parser.add_argument("--fps", type=int, default=24, help="FPS of the sampled video") + parser.add_argument("--height", type=int, default=704, help="Height of image to sample") + parser.add_argument("--width", type=int, default=1280, help="Width of image to sample") + parser.add_argument("--seed", type=int, default=1, help="Random seed") + parser.add_argument("--num_devices", type=int, default=1, help="Number of devices for inference") + parser.add_argument("--cp_size", type=int, default=1, help="Number of cp ranks for multi-gpu inference.") + parser.add_argument("--num_steps", type=float, default=35, help="Number of diffusion sampling steps") + parser.add_argument("--num_video_frames", type=int, default=121, help="Number of video frames to sample") + parser.add_argument( + "--tokenizer_model", type=str, default="nvidia/Cosmos-1.0-Tokenizer-CV8x8x8", help="Mode of video tokenizer" + ) + parser.add_argument("--tokenizer_dir", type=str, default="", help="Directory for video tokenizer") + parser.add_argument("--cosmos_assets_dir", type=str, default="", help="Directory containing cosmos assets") + parser.add_argument("--guardrail_dir", type=str, default="", help="Guardrails weights directory") + parser.add_argument("--nemo_checkpoint", type=str, default="", help="Video diffusion model nemo weights") + parser.add_argument("--t5_cache_dir", type=str, default=None, help="Path to T5 model") + args = parser.parse_args() + return args + + +def print_rank_0(string: str): + rank = torch.distributed.get_rank() + if rank == 0: + print(string) + + +@torch.no_grad() +def encode_for_batch(tokenizer: T5TokenizerFast, encoder: T5EncoderModel, prompts: list[str], max_length: int = 512): + """ + Encode a batch of text prompts to a batch of T5 embeddings. + Parameters: + tokenizer: T5 embedding tokenizer. + encoder: T5 embedding text encoder. + prompts: A batch of text prompts. + max_length: Sequence length of text embedding (defaults to 512). + """ + + batch_encoding = tokenizer.batch_encode_plus( + prompts, + return_tensors="pt", + truncation=True, + padding="max_length", + max_length=max_length, + return_length=True, + return_offsets_mapping=False, + ) + + # We expect all the processing is done on GPU. + input_ids = batch_encoding.input_ids.cuda() + attn_mask = batch_encoding.attention_mask.cuda() + + outputs = encoder(input_ids=input_ids, attention_mask=attn_mask) + encoded_text = outputs.last_hidden_state + + lengths = attn_mask.sum(dim=1).cpu() + for batch_id in range(encoded_text.shape[0]): + encoded_text[batch_id][lengths[batch_id] :] = 0 + + return encoded_text + + +class PosID3D: + def __init__(self, *, max_t=32, max_h=128, max_w=128): + self.max_t = max_t + self.max_h = max_h + self.max_w = max_w + self.generate_pos_id() + + def generate_pos_id(self): + self.grid = torch.stack( + torch.meshgrid( + torch.arange(self.max_t, device="cpu"), + torch.arange(self.max_h, device="cpu"), + torch.arange(self.max_w, device="cpu"), + ), + dim=-1, + ) + + def get_pos_id_3d(self, *, t, h, w): + if t > self.max_t or h > self.max_h or w > self.max_w: + self.max_t = max(self.max_t, t) + self.max_h = max(self.max_h, h) + self.max_w = max(self.max_w, w) + self.generate_pos_id() + return self.grid[:t, :h, :w] + + +def prepare_data_batch(args, t5_embeding_max_length=512): + tokenizer = T5TokenizerFast.from_pretrained("google-t5/t5-11b", cache_dir=args.t5_cache_dir) + text_encoder = T5EncoderModel.from_pretrained("google-t5/t5-11b", cache_dir=args.t5_cache_dir) + text_encoder.to("cuda") + text_encoder.eval() + + print("[args.prompt]: ", args.prompt) + # Encode text to T5 embedding + out = encode_for_batch(tokenizer, text_encoder, args.prompt.split(",")) + encoded_text = torch.tensor(out, dtype=torch.bfloat16) + B, L, C = encoded_text.shape + t5_embed = torch.zeros(B, t5_embeding_max_length, C, dtype=torch.bfloat16) + t5_embed[:, :L, :] = encoded_text + neg_t5_embed = None + t, h, w = args.num_video_frames, args.height, args.width + pt, ph, pw = 1, 2, 2 + state_shape = [ + B, # batch dimension + ((h // 8) // ph) + * ((w // 8) // pw) + * 1, # number of tokens: (h //8) * (w // 8) * 1 -> ((h // 8) // ph) * ((w // 8) // pw) * 1 + 16 * (ph * pw * pt), # token hidden size (channel * patch_spatial * patch_spatial * patch_temporal) + ] + # prepare pos_emb + pos_id_3d = PosID3D() + pt, ph, pw = 1, 2, 2 + pos_ids = rearrange( + # pos_id_3d.get_pos_id_3d(t=t // 4, h=h // 8, w=w // 8), + pos_id_3d.get_pos_id_3d(t=1, h=(h // 8) // ph, w=(w // 8) // pw), + "T H W d -> T (H W) d", + ) + data_batch = { + "video": torch.zeros((B, 3, t, h, w), dtype=torch.uint8).cuda(), + "context_embeddings": t5_embed, + "context_mask": torch.ones(B, t5_embeding_max_length, dtype=torch.bfloat16).cuda(), + "image_size": torch.tensor( + [[args.height, args.width, args.height, args.width]] * B, dtype=torch.bfloat16 + ).cuda(), + "fps": torch.tensor([[args.fps]] * B, dtype=torch.bfloat16).cuda(), + "num_frames": torch.tensor([[args.num_video_frames]] * B, dtype=torch.bfloat16).cuda(), + "padding_mask": torch.zeros((B, 1, args.height, args.width), dtype=torch.bfloat16).cuda(), + "pos_ids": pos_ids, + "latent_shape": [16, t // pt, h // 8 // ph, w // 8 // pw], + } + return data_batch, state_shape + + +def setup_diffusion_pipeline(args): + """ + Initialize DiT model, parallel strategy, and diffusion pipeline for inference. + """ + model_config = DiTModelProvider() + model = model_config.provide() + model = model.cuda().to(torch.bfloat16) + diffusion_pipeline = EDMPipeline(seed=args.seed) + diffusion_pipeline.net = model + return model, diffusion_pipeline, model_config + + +def data_preprocess(data_batch, state_shape): + from dfm.src.megatron.model.dit.dit_data_process import encode_seq_length + + data_batch = {k: v.cuda() if torch.is_tensor(v) else v for k, v in data_batch.items()} + data_batch["inference_fwd"] = True + + data_batch["seq_len_q"] = torch.tensor([state_shape[1]] * state_shape[0]).cuda() + data_batch["seq_len_kv"] = torch.tensor([data_batch["context_embeddings"].shape[1]] * state_shape[0]).cuda() + data_batch = encode_seq_length(data_batch, format="sbhd") + return data_batch + + +def main(args): + # Initialize distributed environment and model parallel groups + Utils.initialize_distributed(1, 1, context_parallel_size=args.cp_size) + model_parallel_cuda_manual_seed(args.seed) + + # Setup model / diffusion pipeline + print_rank_0("setting up diffusion pipeline...") + import random + + model_parallel_cuda_manual_seed(args.seed) + + def set_seed(seed): + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + np.random.seed(seed) + random.seed(seed) + # For deterministic behavior + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + + set_seed(42) + model, diffusion_pipeline, model_config = setup_diffusion_pipeline(args) + + # TODO: load model from checkpoint path argument + new_state = {} + print("loading model....") + state = torch.load("model.pth") + for key, value in state.items(): + if "extra_state" in key: + continue + new_state[key.replace("0.module.", "")] = value + model.load_state_dict(new_state, strict=False) + + print_rank_0("preparing data batch...") + data_batch, state_shape = prepare_data_batch(args) + vae = CausalVideoTokenizer.from_pretrained("Cosmos-0.1-Tokenizer-CV4x8x8") + vae.to("cuda") + + print_rank_0("generating video...") + data_batch = data_preprocess(data_batch, state_shape) + C, T, H, W = data_batch["latent_shape"] + latent = diffusion_pipeline.generate_samples_from_batch( + data_batch=data_batch, + model=model, + guidance=args.guidance, + state_shape=state_shape, + num_steps=args.num_steps, + is_negative_prompt=True if "neg_t5_text_embeddings" in data_batch else False, + ) + rank = torch.distributed.get_rank() + latent = latent[0, None, : state_shape[1]] + latent = rearrange( + latent, + "b (T H W) (ph pw pt c) -> b c (T pt) (H ph) (W pw)", + ph=model_config.patch_spatial, + pw=model_config.patch_spatial, + pt=model_config.patch_temporal, + c=C, + T=T, + H=H, + W=W, + ) + decoded_video = (1.0 + vae.decode(latent / model_config.sigma_data)).clamp(0, 2) / 2 + decoded_video = (decoded_video * 255).to(torch.uint8).permute(0, 2, 3, 4, 1).cpu().numpy() + for i in range(len(decoded_video)): + save_video( + grid=decoded_video[i], + fps=args.fps, + H=args.height, + W=args.width, + video_save_quality=5, + video_save_path=f"idx={i}_rank={rank}_" + args.video_save_path, + ) + print_rank_0(f"saved video to idx={i}_rank={rank}_{args.video_save_path}") + + +if __name__ == "__main__": + args = parse_args() + main(args) diff --git a/examples/megatron/recipes/dit/prepare_energon_dataset_butterfly.py b/examples/megatron/recipes/dit/prepare_energon_dataset_butterfly.py new file mode 100644 index 00000000..f05bdfc8 --- /dev/null +++ b/examples/megatron/recipes/dit/prepare_energon_dataset_butterfly.py @@ -0,0 +1,303 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. 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. + +import os +import pickle +from typing import Callable + +import nemo_run as run +import numpy as np +import pandas as pd +import torch +import torch.distributed as dist +import webdataset as wds +from einops import rearrange +from nemo.collections.common.video_tokenizers.cosmos_tokenizer import CausalVideoTokenizer +from nemo.collections.common.video_tokenizers.utils import read_image, resize_video +from transformers import T5EncoderModel, T5TokenizerFast + + +# examples/megatron/receipes/dit +def initialize_text_encoder(t5_cache_dir): + """ + Initializes the T5 tokenizer and encoder model, loading them from a specified cache directory. + + Args: + t5_cache_dir (str): Path to the cache directory for storing the pretrained model files. + + Returns: + tuple: A tuple containing the tokenizer and encoder model instances. + """ + + # Load tokenizer and text encoder, save in cache directory + tokenizer = T5TokenizerFast.from_pretrained("google-t5/t5-11b", cache_dir=t5_cache_dir) + text_encoder = T5EncoderModel.from_pretrained("google-t5/t5-11b", cache_dir=t5_cache_dir) + text_encoder.to("cuda") + text_encoder.eval() + + return tokenizer, text_encoder + + +# Load dataset from HuggingFace +df = pd.read_parquet("hf://datasets/huggan/smithsonian_butterflies_subset/data/train-00000-of-00001.parquet") +# Load Cosmos tokenizer from HuggingFace + +autoencoder = CausalVideoTokenizer.from_pretrained("Cosmos-0.1-Tokenizer-CV4x8x8") + +# Load T5-XXL text encoder +t5_cache_dir = "" # Use your own custom cache path +tokenizer, text_encoder = initialize_text_encoder(t5_cache_dir) + + +class EncodedSample: + """ + A class representing an encoded sample, containing the text encoding, length, + attention mask, and offset mappings. + + Attributes: + encoded_text (np.ndarray): Encoded text array. + length (int): Length of the encoding. + attn_mask (np.ndarray): Attention mask for the encoding. + offset_mappings (np.ndarray): Mappings for offset positions. + """ + + def __init__(self, encoded_text: np.ndarray, length: int, attn_mask: np.ndarray, offset_mappings: np.ndarray): + self.encoded_text = encoded_text + self.length = length + self.attn_mask = attn_mask + self.offset_mappings = offset_mappings + + def truncate(self) -> None: + """ + Truncates the encoded text, attention mask, and offset mappings to the specified length. + """ + self.encoded_text = self.encoded_text[0 : self.length].astype(np.float16) + self.attn_mask = self.attn_mask[0 : self.length].astype(np.int32) + if self.offset_mappings is not None: + self.offset_mappings = self.offset_mappings[0 : self.length].astype(np.int32) + + +@torch.no_grad() +def encode_for_batch( + tokenizer, encoder, prompts: list[str], truncate: bool = True, max_length=512, output_mapping=True +): + """ + Encodes a batch of text prompts into T5 embeddings. + + Args: + tokenizer: Tokenizer instance for encoding. + encoder: T5 encoder model instance. + prompts (list[str]): List of text prompts to encode. + truncate (bool): If True, truncates the output embeddings. + max_length (int): Maximum length for each encoded prompt. + output_mapping (bool): If True, returns offset mappings for each prompt. + + Returns: + list[EncodedSample]: A list of encoded samples containing text encodings and masks. + """ + batch_encoding = tokenizer.batch_encode_plus( + prompts, + return_tensors="pt", + truncation=True, + padding="max_length", + max_length=max_length, + return_length=True, + return_offsets_mapping=output_mapping, + ) + + # We expect all the processing is done in GPU. + input_ids = batch_encoding.input_ids.cuda() + attn_mask = batch_encoding.attention_mask.cuda() + if output_mapping: + offsets_mapping = batch_encoding["offset_mapping"] + offsets_mapping = offsets_mapping.cpu().numpy() + else: + offsets_mapping = None + + outputs = encoder(input_ids=input_ids, attention_mask=attn_mask) # type: ignore + encoded_text = outputs.last_hidden_state + + lengths = attn_mask.sum(dim=1).cpu() + for batch_id in range(encoded_text.shape[0]): + encoded_text[batch_id][lengths[batch_id] :] = 0 + + encoded_text = encoded_text.cpu().numpy() + attn_mask = attn_mask.cpu().numpy() + + encoded_text = encoded_text[:, :max_length] + attn_mask = attn_mask[:, :max_length] + + out = [] + for idx in range(encoded_text.shape[0]): + if output_mapping: + offsets = offsets_mapping[idx] + else: + offsets = None + + out.append(EncodedSample(encoded_text[idx].astype(np.float16), lengths[idx], attn_mask[idx], offsets)) + if truncate: + for x in out: + x.truncate() + return out + + +def generate_t5_embed(tokenizer, text_encoder, prompt, t5_embeding_max_length=512): + """ + Generates a T5 embedding for a single text prompt. + + Args: + tokenizer: T5 tokenizer instance. + text_encoder: T5 encoder model instance. + prompt (str): The text prompt to encode. + t5_embeding_max_length (int): Maximum length for the embedding. + + Returns: + torch.Tensor: Padded T5 embedding tensor. + """ + # encode text to t5 embedding + out = encode_for_batch(tokenizer, text_encoder, [prompt])[0] + encoded_text = torch.tensor(out.encoded_text, dtype=torch.bfloat16) + + # padding t5 embedding to t5_embeding_max_length + L, C = encoded_text.shape + t5_embed = torch.zeros(1, t5_embeding_max_length, C, dtype=torch.bfloat16) + t5_embed[0, :L] = encoded_text + + return t5_embed + + +def get_start_end_idx_for_this_rank(dataset_size, rank, world_size): + """ + Calculates the start and end indices for distributed processing based on rank. + + Args: + dataset_size (int): Total dataset size. + rank (int): Current process rank. + world_size (int): Total number of processes. + + Returns: + tuple: (start index, end index) for the rank. + """ + split_size = dataset_size // world_size + start_idx = rank * split_size + # The last rank takes the remainder + end_idx = start_idx + split_size if rank != world_size - 1 else dataset_size + return start_idx, end_idx + + +def butterfly_process_func(index, rank): + """ + Generates a sample dictionary with image latent tensor, caption, and metadata. + + Args: + index (int): Index of the dataset row. + rank (int): Current process rank for GPU device selection. + + Returns: + dict: Dictionary containing processed image latents, embeddings, and metadata. + """ + # Access the data from the dataframe + row = df.iloc[index] + image_url = row["image_url"] + image_caption = row["name"] + + # Process image + video = read_image(image_url) + video = rearrange(video, "h w (t c) -> t h w c", t=1) + + # import pdb; pdb.set_trace() + video = resize_video(video, short_size=512) + import mediapy as media + + # Ensure that h and w are divisible by 16 + h, w = video.shape[1:3] + video = media.resize_video(video, shape=(h // 16 * 16, w // 16 * 16)) + batch_video = video[np.newaxis, ...] + + # Bx3xTxHxW + batch_video = rearrange(batch_video, "b t h w c -> b c t h w") + # make video -1...1. Currenlty it has 0-255 + batch_video = (batch_video / 255.0) * 2 - 1 + # Run autoencoder to get latents + + # import pdb; pdb.set_trace() + image_latent = autoencoder.encode(torch.from_numpy(batch_video).to(torch.bfloat16).cuda(device=rank))[0] + image_latent = image_latent.cpu() + + text_embedding = generate_t5_embed(tokenizer, text_encoder, image_caption) + + # Construct sample dictionary + sample = { + "__key__": f"{index:06}", + ".pth": image_latent.to(dtype=torch.bfloat16), + ".pickle": pickle.dumps(text_embedding), + ".json": { + "image_height": batch_video.shape[2], + "image_width": batch_video.shape[3], + # Add additional score as metadata + }, + } + return sample + + +@torch.no_grad() +@run.cli.entrypoint +def prepare(process_func: Callable, output_dir: str = "output_butterfly"): + """ + Prepares a WebDataset using the specified processing function, for distributed settings. + + Args: + process_func (Callable): Function to process each dataset entry. + output_dir (str): Output directory to save processed dataset. + + """ + rank = dist.get_rank() + world_size = torch.distributed.get_world_size() + # rank = 0 + # world_size = 1 + # import pdb; pdb.set_trace() + print(f"Rank {rank} of {world_size} processing {len(df)} samples") + start_idx, end_idx = get_start_end_idx_for_this_rank(len(df), rank, world_size) + print(f"Rank {rank} of {world_size} processing {end_idx - start_idx} samples, from {start_idx} to {end_idx}") + os.makedirs(output_dir, exist_ok=True) + output_tar = os.path.join(output_dir, f"rank{rank}-%06d.tar") + + with wds.ShardWriter(output_tar, maxcount=10000) as sink: + # for i in range(start_idx, end_idx): + from tqdm import tqdm + + for i in tqdm(range(start_idx, end_idx)): + # convert to tqdm + sample = process_func(i, rank) + # Write sample to tar file + sink.write(sample) + + +@run.cli.factory(target=prepare) +def prepare_butterfly_dataset() -> run.Partial: + """ + Prepares the butterfly dataset for distributed training. + + Returns: + run.Partial: Partially configured run for WebDataset preparation. + """ + recipe = run.Partial(prepare, process_func=butterfly_process_func, output_dir="butterfly_webdataset") + return recipe + + +if __name__ == "__main__": + dist.init_process_group("nccl") + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + run.cli.main(prepare, default_factory=prepare_butterfly_dataset) diff --git a/examples/megatron/recipes/dit/pretrain_dit_model.py b/examples/megatron/recipes/dit/pretrain_dit_model.py new file mode 100644 index 00000000..ae4dc54d --- /dev/null +++ b/examples/megatron/recipes/dit/pretrain_dit_model.py @@ -0,0 +1,185 @@ +#!/usr/bin/env python3 +# Copyright (c) 2025, NVIDIA CORPORATION. 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. + +""" +Llama3 8B Pretraining Script with YAML and CLI Configuration Overrides. + +This script provides a flexible way to pretrain Llama3 8B models using Megatron-Bridge with support for +both YAML configuration files and command-line overrides using Hydra-style syntax. + +Examples: + Basic usage with default configuration: + $ torchrun --nproc_per_node=8 pretrain_llama3_8b.py + + Using a custom YAML config file: + $ torchrun --nproc_per_node=8 pretrain_llama3_8b.py --config-file my_custom_config.yaml + + Using CLI overrides only: + $ torchrun --nproc_per_node=8 pretrain_llama3_8b.py model.tensor_model_parallel_size=4 train.train_iters=100000 + + Combining YAML and CLI overrides (CLI takes precedence): + $ torchrun --nproc_per_node=8 pretrain_llama3_8b.py --config-file conf/my_config.yaml \ + model.pipeline_dtype=torch.float16 \ + train.global_batch_size=512 + +Configuration Precedence: + 1. Base configuration from pretrain_config() recipe + 2. YAML overrides from --config-file (if provided) + 3. CLI overrides (highest precedence) + +Supported Override Syntax: + - Standard assignment: key=value + - Nested assignment: section.subsection.key=value + - Addition: +new_key=value + - Deletion: ~key_to_remove + - Type conversion: Automatic for basic types (int, float, bool, str) + - Complex types: torch.dtype, enums, etc. are supported +""" + +import argparse +import logging +import os +import sys +from pathlib import Path +from typing import Tuple + +from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.training.pretrain import pretrain +from megatron.bridge.training.utils.omegaconf_utils import ( + apply_overrides, + create_omegaconf_dict_config, + parse_hydra_overrides, +) +from megatron.bridge.utils.common_utils import get_rank_safe +from omegaconf import OmegaConf + +from dfm.src.megatron.model.dit.dit_step import DITForwardStep +from dfm.src.megatron.recipes.dit.dit import pretrain_config + + +logger: logging.Logger = logging.getLogger(__name__) + + +# Define paths relative to this script's location +# Assumes this script (pretrain_llama3_8b.py) is in Megatron-Bridge/examples/recipes/llama/ +# and the config is in a 'conf' subdirectory. +SCRIPT_DIR: Path = Path(__file__).parent.resolve() +DEFAULT_CONFIG_FILENAME: str = "llama3_8b_pretrain_override_example.yaml" +DEFAULT_CONFIG_FILE_PATH: Path = SCRIPT_DIR / "conf" / DEFAULT_CONFIG_FILENAME + + +def parse_cli_args() -> Tuple[argparse.Namespace, list[str]]: + """Parse command line arguments, separating known script args from OmegaConf overrides.""" + parser = argparse.ArgumentParser( + description="Pretrain Llama3 8B model using Megatron-Bridge with YAML and CLI overrides", + formatter_class=argparse.RawTextHelpFormatter, + ) + parser.add_argument( + "--config-file", + type=str, + default=None, + help="Path to the YAML OmegaConf override file. Default: conf/llama3_8b_pretrain_override_example.yaml", + ) + parser.add_argument("--debug", action="store_true", help="Enable debug logging") + parser.add_argument( + "--dataset-path", + type=str, + default="/opt/VFM/butterfly_webdataset", + help="Path to the dataset directory. Default: /opt/VFM/butterfly_webdataset", + ) + + # Parse known args for the script, remaining will be treated as overrides + args, cli_dotlist_overrides = parser.parse_known_args() + return args, cli_dotlist_overrides + + +def main() -> None: + """ + Entry point for the Llama3 8B pretraining script. + + This function orchestrates the complete configuration workflow: + 1. Loads the base configuration from pretrain_config() recipe + 2. Applies YAML overrides from --config-file (if exists) + 3. Applies CLI overrides using Hydra-style syntax + 4. Starts Megatron pretraining with the final merged configuration + + Configuration merging preserves callable fields (like activation functions) + and handles type conversions automatically. + + Examples of CLI usage: + # Use default config with custom learning rate + torchrun --nproc_per_node=8 pretrain_llama3_8b.py optimizer.lr=0.0002 + + # Custom config file with additional overrides + torchrun --nproc_per_node=8 pretrain_llama3_8b.py --config-file my_config.yaml train.train_iters=50000 + + # Multiple overrides for distributed training + torchrun --nproc_per_node=8 pretrain_llama3_8b.py \ + model.tensor_model_parallel_size=4 \ + model.pipeline_model_parallel_size=2 \ + train.global_batch_size=512 + """ + args, cli_overrides = parse_cli_args() + + print(args) + + logger.info("Megatron-Bridge Llama3 8B Pretraining Script with YAML & CLI Overrides") + logger.info("------------------------------------------------------------------") + + # Load base configuration from the recipe as a Python dataclass + cfg: ConfigContainer = pretrain_config(dataset_path=args.dataset_path) + logger.info("Loaded base configuration") + + # Print configuration on rank 0 + if get_rank_safe() == 0: + cfg.print_yaml() + + # Convert the initial Python dataclass to an OmegaConf DictConfig for merging + merged_omega_conf, excluded_fields = create_omegaconf_dict_config(cfg) + + # Load and merge YAML overrides if a config file is provided + if args.config_file: + logger.debug(f"Loading YAML overrides from: {args.config_file}") + if not os.path.exists(args.config_file): + logger.error(f"Override YAML file not found: {args.config_file}") + sys.exit(1) + yaml_overrides_omega = OmegaConf.load(args.config_file) + merged_omega_conf = OmegaConf.merge(merged_omega_conf, yaml_overrides_omega) + logger.debug("YAML overrides merged successfully.") + + if cli_overrides: + logger.debug(f"Applying Hydra-style command-line overrides: {cli_overrides}") + merged_omega_conf = parse_hydra_overrides(merged_omega_conf, cli_overrides) + logger.debug("Hydra-style command-line overrides applied successfully.") + + # # Apply the final merged OmegaConf configuration back to the original ConfigContainer + logger.debug("Applying final merged configuration back to Python ConfigContainer...") + final_overrides_as_dict = OmegaConf.to_container(merged_omega_conf, resolve=True) + # Apply overrides while preserving excluded fields + apply_overrides(cfg, final_overrides_as_dict, excluded_fields) + + # Display final configuration + if get_rank_safe() == 0: + logger.info("--- Final Merged Configuration ---") + cfg.print_yaml() + logger.info("----------------------------------") + + # Start training + logger.debug("Starting pretraining...") + pretrain(config=cfg, forward_step_func=DITForwardStep()) + + +if __name__ == "__main__": + main()