nicolabs/Open-Sora
1
0
Fork 0
mirror of https://github.com/hpcaitech/Open-Sora.git synced 2026-05-21 11:59:01 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

added sp for stdit3 (#131)

Browse source
This commit is contained in:
Frank Lee 2024-06-14 11:10:43 +08:00 committed by GitHub
parent 10e1e1062d
commit a887f54711
6 changed files with 152 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
opensora/models/layers/blocks.py
View file

@ -360,7 +360,6 @@ class SeqParallelAttention(Attention):
norm_layer: nn.Module = LlamaRMSNorm,
enable_flash_attn: bool = False,
rope=None,
qk_norm_legacy: bool = False,
) -> None:
assert rope is None, "Rope is not supported in SeqParallelAttention"
super().__init__(
@ -378,7 +377,6 @@ class SeqParallelAttention(Attention):
B, N, C = x.shape # for sequence parallel here, the N is a local sequence length
qkv = self.qkv(x)
qkv_shape = (B, N, 3, self.num_heads, self.head_dim)
qkv = qkv.view(qkv_shape)
sp_group = get_sequence_parallel_group()
@ -496,21 +494,19 @@ class SeqParallelMultiHeadCrossAttention(MultiHeadCrossAttention):
# query/value: img tokens; key: condition; mask: if padding tokens
sp_group = get_sequence_parallel_group()
sp_size = dist.get_world_size(sp_group)
B, SUB_N, C = x.shape
B, SUB_N, C = x.shape # [B, TS/p, C]
N = SUB_N * sp_size
# shape:
# q, k, v: [B, SUB_N, NUM_HEADS, HEAD_DIM]
q = self.q_linear(x).view(B, -1, self.num_heads, self.head_dim)
kv = self.kv_linear(cond).view(B, -1, 2, self.num_heads, self.head_dim)
q = self.q_linear(x).view(1, -1, self.num_heads, self.head_dim)
kv = self.kv_linear(cond).view(1, -1, 2, self.num_heads, self.head_dim)
kv = split_forward_gather_backward(kv, get_sequence_parallel_group(), dim=3, grad_scale="down")
k, v = kv.unbind(2)
# apply all_to_all to gather sequence and split attention heads
q = all_to_all(q, sp_group, scatter_dim=2, gather_dim=1)
k = split_forward_gather_backward(k, get_sequence_parallel_group(), dim=2, grad_scale="down")
v = split_forward_gather_backward(v, get_sequence_parallel_group(), dim=2, grad_scale="down")
q = q.view(1, -1, self.num_heads // sp_size, self.head_dim)
k = k.view(1, -1, self.num_heads // sp_size, self.head_dim)
v = v.view(1, -1, self.num_heads // sp_size, self.head_dim)

38
opensora/models/stdit/stdit3.py
View file

@ -1,14 +1,18 @@
import numpy as np
import torch
import torch.nn as nn
import torch.distributed as dist
from einops import rearrange
from rotary_embedding_torch import RotaryEmbedding
from timm.models.layers import DropPath
from timm.models.vision_transformer import Mlp
from transformers import PretrainedConfig, PreTrainedModel
from opensora.acceleration.communications import gather_forward_split_backward, split_forward_gather_backward
from opensora.acceleration.parallel_states import get_sequence_parallel_group
from opensora.acceleration.checkpoint import auto_grad_checkpoint
from opensora.models.layers.blocks import (
SeqParallelMultiHeadCrossAttention,
SeqParallelAttention,
Attention,
CaptionEmbedder,
MultiHeadCrossAttention,
@ -43,12 +47,15 @@ class STDiT3Block(nn.Module):
self.temporal = temporal
self.hidden_size = hidden_size
self.enable_flash_attn = enable_flash_attn
self._enable_sequence_parallelism = enable_sequence_parallelism
assert not enable_sequence_parallelism, "Sequence parallelism is not supported in STDiT3Block"
attn_cls = Attention
mha_cls = MultiHeadCrossAttention
self.enable_sequence_parallelism = enable_sequence_parallelism
if self.enable_sequence_parallelism and not temporal:
attn_cls = SeqParallelAttention
mha_cls = SeqParallelMultiHeadCrossAttention
else:
attn_cls = Attention
mha_cls = MultiHeadCrossAttention
self.norm1 = get_layernorm(hidden_size, eps=1e-6, affine=False, use_kernel=enable_layernorm_kernel)
self.attn = attn_cls(
hidden_size,
@ -215,6 +222,7 @@ class STDiT3(PreTrainedModel):
self.drop_path = config.drop_path
self.enable_flash_attn = config.enable_flash_attn
self.enable_layernorm_kernel = config.enable_layernorm_kernel
self.enable_sequence_parallelism = config.enable_sequence_parallelism
# input size related
self.patch_size = config.patch_size
@ -380,17 +388,29 @@ class STDiT3(PreTrainedModel):
x = self.x_embedder(x) # [B, N, C]
x = rearrange(x, "B (T S) C -> B T S C", T=T, S=S)
x = x + pos_emb
# shard over the sequence dim if sp is enabled
if self.enable_sequence_parallelism:
x = split_forward_gather_backward(x, get_sequence_parallel_group(), dim=2, grad_scale="down")
S = S // dist.get_world_size(get_sequence_parallel_group())
x = rearrange(x, "B T S C -> B (T S) C", T=T, S=S)
# === blocks ===
for spatial_block, temporal_block in zip(self.spatial_blocks, self.temporal_blocks):
x = auto_grad_checkpoint(spatial_block, x, y, t_mlp, y_lens, x_mask, t0_mlp, T, S)
x = auto_grad_checkpoint(temporal_block, x, y, t_mlp, y_lens, x_mask, t0_mlp, T, S)
if self.enable_sequence_parallelism:
x = rearrange(x, "B (T S) C -> B T S C", T=T, S=S)
x = gather_forward_split_backward(x, get_sequence_parallel_group(), dim=2, grad_scale="up")
S = S * dist.get_world_size(get_sequence_parallel_group())
x = rearrange(x, "B T S C -> B (T S) C", T=T, S=S)
# === final layer ===
x = self.final_layer(x, t, x_mask, t0, T, S)
x = self.unpatchify(x, T, H, W, Tx, Hx, Wx)
# cast to float32 for better accuracy
x = x.to(torch.float32)
return x

2
opensora/utils/train_utils.py
View file

@ -14,6 +14,7 @@ from .misc import get_logger
def create_colossalai_plugin(plugin, dtype, grad_clip, sp_size):
if plugin == "zero2":
assert sp_size == 1, "Zero2 plugin does not support sequence parallelism"
plugin = LowLevelZeroPlugin(
stage=2,
precision=dtype,
@ -22,6 +23,7 @@ def create_colossalai_plugin(plugin, dtype, grad_clip, sp_size):
)
set_data_parallel_group(dist.group.WORLD)
elif plugin == "zero2-seq":
assert sp_size > 1, "Zero2-seq plugin requires sequence parallelism"
plugin = ZeroSeqParallelPlugin(
sp_size=sp_size,
stage=2,

1
scripts/train.py
View file

@ -140,6 +140,7 @@ def main():
in_channels=vae_out_channels,
caption_channels=text_encoder_output_dim,
model_max_length=text_encoder_model_max_length,
enable_sequence_parallelism=cfg.get("sp_size", 1) > 1
)
.to(device, dtype)
.train()

16
tests/test_lr_scheduler.py
View file

@ -1,25 +1,25 @@
import torch
from torch.optim import Adam
from torchvision.models import resnet50
from tqdm import tqdm
from opensora.utils.lr_scheduler import LinearWarmupLR
def test_lr_scheduler():
warmup_steps = 200
model = resnet50().cuda()
optimizer = Adam(model.parameters(), lr=0.01)
scheduler = LinearWarmupLR(optimizer, warmup_steps=10)
scheduler = LinearWarmupLR(optimizer, warmup_steps=warmup_steps)
current_lr = scheduler.get_lr()[0]
data = torch.rand(128, 3, 224, 224).cuda()
for i in range(100):
data = torch.rand(1, 3, 224, 224).cuda()
for i in tqdm(range(warmup_steps*2)):
out = model(data)
out.mean().backward()
optimizer.step()
scheduler.step()
if i >= 10:
if i >= warmup_steps:
assert scheduler.get_lr()[0] == 0.01
else:
assert scheduler.get_lr()[0] > current_lr, f"{scheduler.get_lr()[0]} <= {current_lr}"

108
tests/test_stdit3_sequence_parallelism.py Normal file
View file

@ -0,0 +1,108 @@
import colossalai
import torch
import torch.distributed as dist
from opensora.models.stdit.stdit3 import STDiT3Config, STDiT3
from colossalai.testing import spawn, free_port
from opensora.acceleration.parallel_states import set_data_parallel_group, set_sequence_parallel_group
from colossalai.utils.common import set_seed
def get_sample_data():
x = torch.rand([1, 4, 15, 20, 27], dtype=torch.bfloat16) # (B, C, T, H, W)
timestep = torch.Tensor([924.]).to(torch.bfloat16)
y = torch.rand(1, 1, 300, 4096, dtype=torch.bfloat16)
mask = torch.ones([1, 300], dtype=torch.int32)
x_mask = torch.ones([1, 15]).bool()
fps = torch.Tensor([25.]).to(torch.bfloat16)
height = torch.Tensor([166.]).to(torch.bfloat16)
width = torch.Tensor([221.]).to(torch.bfloat16)
return dict(x=x, timestep=timestep, y=y, mask=mask, x_mask=x_mask, fps=fps, height=height, width=width)
def get_stdit3_config(enable_sequence_parallelism = False):
config = {
"caption_channels": 4096,
"class_dropout_prob": 0.0,
"depth": 1,
"drop_path": 0.0,
"enable_flash_attn": True,
"enable_layernorm_kernel": True,
"enable_sequence_parallelism": enable_sequence_parallelism,
"freeze_y_embedder": True,
"hidden_size": 1152,
"in_channels": 4,
"input_size": [
None,
None,
None
],
"input_sq_size": 512,
"mlp_ratio": 4.0,
"model_max_length": 300,
"model_type": "STDiT3",
"num_heads": 16,
"only_train_temporal": False,
"patch_size": [
1,
2,
2
],
"pred_sigma": True,
"qk_norm": True,
"skip_y_embedder": False,
}
return STDiT3Config(**config)
def run_model(rank, world_size, port):
colossalai.launch({}, rank=rank, world_size=world_size, port=port, host="localhost")
# prepare data
data = get_sample_data()
data = {
k: v.cuda()
for k, v in data.items()
}
# test single-gpu outptu
set_seed(1024)
non_dist_model_cfg = get_stdit3_config(enable_sequence_parallelism=False)
non_dist_model = STDiT3(non_dist_model_cfg).cuda().to(torch.bfloat16)
non_dist_out = non_dist_model(**data)
non_dist_out.mean().backward()
# run seq parallelism
set_sequence_parallel_group(dist.group.WORLD)
set_seed(1024)
dist_model_cfg = get_stdit3_config(enable_sequence_parallelism=True)
dist_model = STDiT3(dist_model_cfg).cuda().to(torch.bfloat16)
dist_out = dist_model(**data)
dist_out.mean().backward()
# run all reduce for gradients
for param in dist_model.parameters():
if param.grad is not None:
dist.all_reduce(param.grad, group=dist.group.WORLD)
param.grad /= world_size
# ensure model weights are equal
for (p1, p2) in zip(non_dist_model.parameters(), dist_model.parameters()):
assert torch.equal(p1, p2)
# check
torch.testing.assert_close(non_dist_out, dist_out)
for ((n1, p1), (n2, p2)) in zip(non_dist_model.named_parameters(), dist_model.named_parameters()):
assert n1 == n2
if p1.grad is not None and p2.grad is not None:
if not torch.allclose(p1.grad, p2.grad,rtol=1e-2, atol=1e-4):
if dist.get_rank() == 0:
print(f"gradient of {n1} is not equal, {p1.grad} vs {p2.grad}")
else:
assert p1.grad is None and p2.grad is None
def test_stdit3_sp():
spawn(run_model, 2)
if __name__ == '__main__':
test_stdit3_sp()