hotfix accuracy and step num (#27)

* fix

* update settings

* Update dtype to torch.float32 in DiT model

* Remove unnecessary assert statement in DiT class

* Refactor logging and checkpoint saving in train_img.py

Author: oahzxl

opendit/models/dit.py

@@ -15,9 +15,8 @@
 import torch
 import torch.distributed as dist
 import torch.nn as nn
-import torch.nn.functional as F
 import torch.utils.checkpoint
-from timm.models.vision_transformer import Mlp, PatchEmbed, use_fused_attn
+from timm.models.vision_transformer import Mlp, PatchEmbed
 from torch.jit import Final
 
 from opendit.models.clip import TextEmbedder
@@ -158,7 +157,6 @@ def __init__(
         self.num_heads = num_heads
         self.head_dim = dim // num_heads
         self.scale = self.head_dim**-0.5
-        self.fused_attn = use_fused_attn()
 
         self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
         self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
@@ -236,13 +234,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
                 dropout_p=self.attn_drop.p if self.training else 0.0,
                 softmax_scale=self.scale,
             )
-        elif self.fused_attn:
-            x = F.scaled_dot_product_attention(
-                q,
-                k,
-                v,
-                dropout_p=self.attn_drop.p if self.training else 0.0,
-            )
         else:
             dtype = q.dtype
             q = q * self.scale
@@ -260,7 +251,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             if self.sequence_parallel_size == 1
             else (B, N * self.sequence_parallel_size, num_heads * self.head_dim)
         )
-        x = x.transpose(1, 2).reshape(x_output_shape)
+        if self.enable_flashattn:
+            x = x.reshape(x_output_shape)
+        else:
+            x = x.transpose(1, 2).reshape(x_output_shape)
+
         if self.sequence_parallel_size > 1:
             # Todo: Use all_to_all_single for x
             # x = x.reshape(1, -1, num_heads * self.head_dim)
@@ -355,6 +350,7 @@ def __init__(
         enable_layernorm_kernel=False,
         enable_modulate_kernel=False,
         sequence_parallel_size=1,
+        dtype=torch.float32,
     ):
         super().__init__()
         self.learn_sigma = learn_sigma
@@ -363,6 +359,12 @@ def __init__(
         self.patch_size = patch_size
         self.num_heads = num_heads
         self.sequence_parallel_size = sequence_parallel_size
+        self.dtype = dtype
+        if enable_flashattn:
+            assert dtype in [
+                torch.float16,
+                torch.bfloat16,
+            ], f"Flash attention only supports float16 and bfloat16, but got {self.dtype}"
 
         self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, hidden_size, bias=True)
         self.t_embedder = TimestepEmbedder(hidden_size)
@@ -470,6 +472,10 @@ def forward(self, x, t, y):
 
         # Todo: Mock video input by repeating the same frame for all timesteps
         # x = torch.randn(2, 256, 1152).to(torch.bfloat16).cuda()
+
+        # origin inputs should be float32, cast to specified dtype
+        x = x.to(self.dtype)
+
         x = self.x_embedder(x) + self.pos_embed  # (N, T, D), where T = H * W / patch_size ** 2
         t = self.t_embedder(t, dtype=x.dtype)  # (N, D)
         y = self.y_embedder(y, self.training)  # (N, D)
@@ -490,6 +496,9 @@ def forward(self, x, t, y):
 
         x = self.final_layer(x, c)  # (N, T, patch_size ** 2 * out_channels)
         x = self.unpatchify(x)  # (N, out_channels, H, W)
+
+        # cast to float32 for better accuracy
+        x = x.to(torch.float32)
         return x
 
     def forward_with_cfg(self, x, t, y, cfg_scale):

sample.py

@@ -34,7 +34,18 @@ def main(args):
 
     # Load model:
     latent_size = args.image_size // 8
-    model = DiT_models[args.model](input_size=latent_size, num_classes=args.num_classes).to(device)
+    dtype = torch.float32
+    model = (
+        DiT_models[args.model](
+            input_size=latent_size,
+            num_classes=args.num_classes,
+            enable_flashattn=False,
+            enable_layernorm_kernel=False,
+            dtype=dtype,
+        )
+        .to(device)
+        .to(dtype)
+    )
     # Auto-download a pre-trained model or load a custom DiT checkpoint from train.py:
     ckpt_path = args.ckpt or f"DiT-XL-2-{args.image_size}x{args.image_size}.pt"
     state_dict = find_model(ckpt_path)

train_img.py

@@ -14,6 +14,7 @@
 
 import colossalai
 import torch
+import torch.distributed as dist
 from colossalai.booster import Booster
 from colossalai.booster.plugin import LowLevelZeroPlugin
 from colossalai.cluster import DistCoordinator
@@ -60,6 +61,7 @@ def main(args):
     model_string_name = args.model.replace("/", "-")
     # Create an experiment folder
     experiment_dir = f"{args.outputs}/{experiment_index:03d}-{model_string_name}"
+    dist.barrier()
     if coordinator.is_master():
         os.makedirs(experiment_dir, exist_ok=True)
         with open(f"{experiment_dir}/config.txt", "w") as f:
@@ -113,6 +115,7 @@ def main(args):
             enable_layernorm_kernel=args.enable_layernorm_kernel,
             enable_modulate_kernel=args.enable_modulate_kernel,
             sequence_parallel_size=args.sequence_parallel_size,
+            dtype=dtype,
         )
         .to(device)
         .to(dtype)
@@ -208,7 +211,6 @@ def main(args):
                 with torch.no_grad():
                     # Map input images to latent space + normalize latents:
                     x = vae.encode(x).latent_dist.sample().mul_(0.18215)
-                    x = x.to(dtype)
 
                 # Diffusion
                 t = torch.randint(0, diffusion.num_timesteps, (x.shape[0],), device=device)
@@ -224,11 +226,15 @@ def main(args):
 
                 # Log loss values:
                 all_reduce_mean(loss)
-                if coordinator.is_master() and (step + 1) % args.log_every == 0:
-                    pbar.set_postfix({"loss": loss.item()})
-                    writer.add_scalar("loss", loss.item(), epoch * num_steps_per_epoch + step)
+                global_step = epoch * num_steps_per_epoch + step
+                pbar.set_postfix({"loss": loss.item(), "step": step, "global_step": global_step})
 
-                if args.ckpt_every > 0 and (step + 1) % args.ckpt_every == 0:
+                # Log to tensorboard
+                if coordinator.is_master() and (global_step + 1) % args.log_every == 0:
+                    writer.add_scalar("loss", loss.item(), global_step)
+
+                # Save checkpoint
+                if args.ckpt_every > 0 and (global_step + 1) % args.ckpt_every == 0:
                     logger.info(f"Saving checkpoint...")
                     save(
                         booster,

train_video.py

@@ -106,6 +106,7 @@ def main(args):
             enable_layernorm_kernel=args.enable_layernorm_kernel,
             enable_modulate_kernel=args.enable_modulate_kernel,
             sequence_parallel_size=args.sequence_parallel_size,
+            dtype=dtype,
         )
         .to(device)
         .to(dtype)