[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

examples/pytorch/custom_handler_fp8_fsdp1n2_compile/train.py

@@ -1,20 +1,21 @@
 import argparse
-from dataclasses import dataclass
 import logging
+from dataclasses import dataclass
 
-import torch.distributed as dist
 import lightning as L
 import torch
+import torch.distributed as dist
 import torch.nn as nn
 import torch.nn.functional as F
-from lightning.pytorch.demos import Transformer, WikiText2
+from lightning.pytorch.demos import WikiText2
 from lightning.pytorch.strategies import FSDPStrategy, ModelParallelStrategy
 from torch.distributed.fsdp import BackwardPrefetch, MixedPrecision
 from torch.utils.data import DataLoader
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 log = logging.getLogger(__name__)
 
+
 @dataclass
 class Args:
     vocab_size: int = 32000
@@ -24,6 +25,7 @@ class Args:
     enable_gradient_checkpointing: bool = False
     enable_fsdp2: bool = False
 
+
 class SimpleLayer(nn.Module):
     def __init__(self, hidden_size):
         super(SimpleLayer, self).__init__()
@@ -37,6 +39,7 @@ def forward(self, x):
         x = self.activation(x)
         return x
 
+
 class InnerModel(nn.Module):
     def __init__(self, num_layers, hidden_size, vocab_size=32000):
         super(InnerModel, self).__init__()
@@ -46,7 +49,6 @@ def __init__(self, num_layers, hidden_size, vocab_size=32000):
         self.layers = nn.ModuleList([SimpleLayer(hidden_size) for _ in range(num_layers)])
         self.lm_head = nn.Linear(hidden_size, vocab_size)
 
-
     def forward(self, x):
         x = self.embedding(x)
         # Pass the input through each layer sequentially
@@ -66,14 +68,15 @@ def forward(self, *args, **kwargs):
 
 
 class LanguageModel(L.LightningModule):
-    def __init__(self, 
-                vocab_size=32000,
-                enable_fp8 = False, 
-                enable_fsdp2 = False,
-                enable_torch_compile = False,
-                enable_gradient_checkpointing = False,
-                enable_cpu_offload = False
-                ):
+    def __init__(
+        self,
+        vocab_size=32000,
+        enable_fp8=False,
+        enable_fsdp2=False,
+        enable_torch_compile=False,
+        enable_gradient_checkpointing=False,
+        enable_cpu_offload=False,
+    ):
         super().__init__()
         self.model = None
         self.vocab_size = vocab_size
@@ -88,10 +91,11 @@ def __init__(self,
         }  # only used for FP8 training
 
     def log_model_stage(self, stage: str):
-        """
-        Logs the current state of the model with a description of the stage.
+        """Logs the current state of the model with a description of the stage.
+
         Args:
             stage (str): Description of the current model stage.
+
         """
         log.warning(f"Model at stage: {stage}\n{self.model}")
 
@@ -129,7 +133,7 @@ def configure_fsdp2(self):
 
     def configure_fp8(self):
         # Setup fp8 training, if enable_fp8 is false, it will create a fake handler
-        from handlers.fp8_training_handler import FP8Config, Float8TrainingHandler
+        from handlers.fp8_training_handler import Float8TrainingHandler, FP8Config
 
         fp8_config = FP8Config(
             enable_fp8=self.enable_fp8,
@@ -207,13 +211,14 @@ def train(args):
     dataset = WikiText2()
     train_dataloader = DataLoader(dataset, num_workers=8, batch_size=1)
 
-    model = LanguageModel(vocab_size=args.vocab_size,
-                          enable_fp8 = args.enable_fp8,
-                          enable_fsdp2 = args.enable_fsdp2,
-                          enable_torch_compile = args.enable_torch_compile,
-                          enable_gradient_checkpointing = args.enable_gradient_checkpointing,
-                          enable_cpu_offload = args.enable_cpu_offload,
-                          )
+    model = LanguageModel(
+        vocab_size=args.vocab_size,
+        enable_fp8=args.enable_fp8,
+        enable_fsdp2=args.enable_fsdp2,
+        enable_torch_compile=args.enable_torch_compile,
+        enable_gradient_checkpointing=args.enable_gradient_checkpointing,
+        enable_cpu_offload=args.enable_cpu_offload,
+    )
 
     if args.enable_fsdp2:
         strategy = ModelParallelStrategy(

examples/pytorch/fp8_fsdp_compile/train.py

@@ -50,7 +50,6 @@ def module_filter_fn(mod: torch.nn.Module, fqn: str):
 
         self.model = torch.compile(model)
 
-
     def training_step(self, batch):
         input, target = batch
         output = self.model(input, target)
@@ -85,4 +84,4 @@ def train():
 if __name__ == "__main__":
     torch.set_float32_matmul_precision("high")
 
-    train()
+    train()