fix lint errors.

Author: sajadn

dfm/src/common/utils/dynamic_import.py

@@ -14,6 +14,7 @@
 
 import importlib
 
+
 def dynamic_import(full_path):
     """
     Dynamically import a class or function from a given full path.

dfm/src/common/utils/save_video.py

@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
 import imageio
 import numpy as np
+import torch
 
 
 def save_video(
@@ -25,7 +25,6 @@ def save_video(
     video_save_quality: int,
     video_save_path: str,
 ):
-
     kwargs = {
         "fps": fps,
         "quality": video_save_quality,
@@ -34,5 +33,14 @@ def save_video(
         "output_params": ["-f", "mp4"],
     }
 
-    print('video_save_path', video_save_path)
-    imageio.mimsave(video_save_path, grid, "mp4", **kwargs)
+    print("video_save_path", video_save_path)
+    imageio.mimsave(video_save_path, grid, "mp4", **kwargs)
+
+
+def print_dict(dict):
+    for key, value in dict.items():
+        if isinstance(value, torch.Tensor):
+            print(key, value.shape)
+        else:
+            print(key, value)
+    print("-" * 40)

dfm/src/common/utils/torch_split_tensor_for_cp.py

@@ -14,7 +14,7 @@
 
 import torch
 from torch import Tensor
-from torch.distributed import ProcessGroup, all_gather, get_process_group_ranks, get_world_size
+from torch.distributed import ProcessGroup, all_gather, get_world_size
 
 
 def cat_outputs_cp(x: Tensor, seq_dim: int, cp_group: ProcessGroup) -> Tensor:

dfm/src/megatron/data/dit/base.py

@@ -12,14 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from copy import deepcopy
+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
-from torch.utils.data import DataLoader
-from typing_extensions import Self
-import logging
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -107,7 +106,7 @@ def __init__(
         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.task_encoder = task_encoder
         self.init_global_step = 0
         self.train_dataloader_object = None
         self.val_dataloader_object = None
@@ -116,8 +115,7 @@ def __init__(
         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'):
+    def datasets_provider(self, worker_config, split: Literal["train", "val"] = "val"):
         """
         Provide the dataset for training or validation.
 
@@ -132,7 +130,7 @@ def datasets_provider(self, worker_config, split: Literal['train', 'val'] = 'val
         Dataset: The dataset configured for the specified split.
         """
 
-        if split not in {'train', 'val'}:
+        if split not in {"train", "val"}:
             raise ValueError("Invalid value for split. Allowed values are 'train' or 'val'.")
 
         if split == "train":
@@ -194,7 +192,7 @@ def train_dataloader(self) -> Any:
                 worker_debug_path=None,
                 worker_log_level=0,
             )
-        train_dataset = self.datasets_provider(worker_config, split='train')
+        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
@@ -232,7 +230,7 @@ def val_dataloader(self):
                 worker_debug_path=None,
                 worker_log_level=0,
             )
-        val_dataset = self.datasets_provider(worker_config, split='val')
+        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
@@ -285,7 +283,7 @@ def state_dict(self) -> Dict[str, Any]:
                 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}
+            return {"dataloader_state": state, "consumed_samples": consumed_samples}
 
         logger.warning("trainer object not connected to data module object returning empty state")
         return {}
@@ -300,14 +298,14 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
         Parameters:
         state_dict (Dict[str, Any]): The state dictionary containing the saved state of the data module.
         """
-        if not 'dataloader_state' in state_dict:
+        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']
+        state = state_dict["dataloader_state"]
         try:
             if self.trainer:
                 self.trainer.datamodule.train_dataloader().restore_state_global(state)
@@ -331,13 +329,11 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
             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']
+        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,
         )
-
-

dfm/src/megatron/data/dit/diffusion_energon_datamodule.py

@@ -14,16 +14,17 @@
 
 # pylint: disable=C0115,C0116,C0301
 
-from dataclasses import dataclass
 import logging
+from dataclasses import dataclass
 from typing import Any, Dict, Literal
 
-from torch import int_repr
-
-from dfm.src.megatron.data.dit.diffusion_taskencoder import BasicDiffusionTaskEncoder
 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.dit.base import EnergonMultiModalDataModule
+from dfm.src.megatron.data.dit.diffusion_taskencoder import BasicDiffusionTaskEncoder
+
 
 @dataclass(kw_only=True)
 class DiffusionDataModuleConfig(DatasetProvider):
@@ -40,18 +41,19 @@ def __post_init__(self):
         self.dataset = DiffusionDataModule(
             path=self.path,
             seq_length=self.seq_length,
-            task_encoder=BasicDiffusionTaskEncoder(seq_length=self.task_encoder_seq_length, packing_buffer_size=self.packing_buffer_size),
+            task_encoder=BasicDiffusionTaskEncoder(
+                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)
+            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):

dfm/src/megatron/data/dit/diffusion_taskencoder.py

@@ -13,17 +13,19 @@
 # limitations under the License.
 
 
+import random
+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 DefaultTaskEncoder, SkipSample
-from megatron.energon.task_encoder.cooking import Cooker, basic_sample_keys
 from megatron.energon.task_encoder.base import stateless
-from dfm.src.megatron.data.dit.sequence_packing_utils import first_fit_decreasing
+from megatron.energon.task_encoder.cooking import Cooker, basic_sample_keys
+
 from dfm.src.megatron.data.dit.dit_sample import DiffusionSample
-from typing import List
-from megatron.core import parallel_state
-import random
+from dfm.src.megatron.data.dit.sequence_packing_utils import first_fit_decreasing
 
 
 def cook(sample: dict) -> dict:
@@ -123,7 +125,7 @@ def encode_sample(self, sample: dict) -> dict:
         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=}')
+            warnings.warn(f"skipping {video_latent.shape=} not divisible by {tpcp_size=}")
             raise SkipSample()
 
         video_latent = rearrange(
@@ -178,10 +180,10 @@ def encode_sample(self, sample: dict) -> dict:
             loss_mask = torch.ones(seq_len, dtype=torch.bfloat16)
 
         return DiffusionSample(
-            __key__=sample['__key__'],
-            __restore_key__=sample['__restore_key__'],
+            __key__=sample["__key__"],
+            __restore_key__=sample["__restore_key__"],
             __subflavor__=None,
-            __subflavors__=sample['__subflavors__'],
+            __subflavors__=sample["__subflavors__"],
             video=video_latent,
             t5_text_embeddings=t5_text_embeddings,
             t5_text_mask=t5_text_mask,
@@ -217,14 +219,14 @@ def cat(attr):
             __restore_key__=(),  # Will be set by energon based on `samples`
             __subflavor__=None,
             __subflavors__=samples[0].__subflavors__,
-            video=cat('video'),
-            t5_text_embeddings=cat('t5_text_embeddings'),
-            t5_text_mask=cat('t5_text_mask'),
-            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'),
+            video=cat("video"),
+            t5_text_embeddings=cat("t5_text_embeddings"),
+            t5_text_mask=cat("t5_text_mask"),
+            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
@@ -247,6 +249,7 @@ def batch(self, samples: List[DiffusionSample]) -> dict:
             latent_shape=sample.latent_shape,
         )
 
+
 class PosID3D:
     def __init__(self, *, max_t=32, max_h=128, max_w=128):
         self.max_t = max_t

dfm/src/megatron/data/dit/dit_sample.py

@@ -13,9 +13,11 @@
 # limitations under the License.
 
 from dataclasses import dataclass
-from typing import Optional, Any
-from megatron.energon import Sample
+from typing import Any, Optional
+
 import torch
+from megatron.energon import Sample
+
 
 @dataclass
 class DiffusionSample(Sample):
@@ -93,4 +95,3 @@ def __lt__(self, other: Any) -> bool:
         elif isinstance(other, int):
             return self.seq_len_q.item() < other
         raise NotImplementedError
-

dfm/src/megatron/data/dit/sequence_packing_utils.py

@@ -14,7 +14,6 @@
 
 
 from typing import List
-import torch
 
 
 def find_first_bin_that_fits(bins: List[List[int]], s: int, bin_size: int) -> int:
@@ -56,6 +55,7 @@ def first_fit(seqlens: List[int], pack_size: int) -> List[List[int]]:
             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.
@@ -72,6 +72,7 @@ def first_fit_decreasing(seqlens: List[int], pack_size: int) -> List[List[int]]:
     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
@@ -102,4 +103,3 @@ def concat_pad(tensor_list, max_seq_length):
         current_index += length
 
     return result
-