model.py 19.3 KB
Newer Older
gushiqiao's avatar
gushiqiao committed
1
import gc
Yang Yong (雍洋)'s avatar
Yang Yong (雍洋) committed
2
import glob
3
4
import os

helloyongyang's avatar
helloyongyang committed
5
import torch
6
import torch.distributed as dist
helloyongyang's avatar
helloyongyang committed
7
import torch.nn.functional as F
PengGao's avatar
PengGao committed
8
9
10
from loguru import logger
from safetensors import safe_open

11
from lightx2v.models.networks.wan.infer.feature_caching.transformer_infer import (
12
13
    WanTransformerInferAdaCaching,
    WanTransformerInferCustomCaching,
Rongjin Yang's avatar
Rongjin Yang committed
14
15
    WanTransformerInferDualBlock,
    WanTransformerInferDynamicBlock,
PengGao's avatar
PengGao committed
16
    WanTransformerInferFirstBlock,
Musisoul's avatar
Musisoul committed
17
    WanTransformerInferMagCaching,
PengGao's avatar
PengGao committed
18
19
20
    WanTransformerInferTaylorCaching,
    WanTransformerInferTeaCaching,
)
21
22
23
from lightx2v.models.networks.wan.infer.offload.transformer_infer import (
    WanOffloadTransformerInfer,
)
PengGao's avatar
PengGao committed
24
25
26
27
28
29
30
31
from lightx2v.models.networks.wan.infer.post_infer import WanPostInfer
from lightx2v.models.networks.wan.infer.pre_infer import WanPreInfer
from lightx2v.models.networks.wan.infer.transformer_infer import (
    WanTransformerInfer,
)
from lightx2v.models.networks.wan.weights.pre_weights import WanPreWeights
from lightx2v.models.networks.wan.weights.transformer_weights import (
    WanTransformerWeights,
32
)
Yang Yong(雍洋)'s avatar
Yang Yong(雍洋) committed
33
from lightx2v.utils.custom_compiler import CompiledMethodsMixin, compiled_method
34
from lightx2v.utils.envs import *
35
from lightx2v.utils.utils import *
helloyongyang's avatar
helloyongyang committed
36

37
38
39
40
41
try:
    import gguf
except ImportError:
    gguf = None

helloyongyang's avatar
helloyongyang committed
42

Yang Yong(雍洋)'s avatar
Yang Yong(雍洋) committed
43
class WanModel(CompiledMethodsMixin):
helloyongyang's avatar
helloyongyang committed
44
45
46
    pre_weight_class = WanPreWeights
    transformer_weight_class = WanTransformerWeights

47
    def __init__(self, model_path, config, device, model_type="wan2.1"):
Yang Yong(雍洋)'s avatar
Yang Yong(雍洋) committed
48
        super().__init__()
helloyongyang's avatar
helloyongyang committed
49
50
        self.model_path = model_path
        self.config = config
51
52
        self.cpu_offload = self.config.get("cpu_offload", False)
        self.offload_granularity = self.config.get("offload_granularity", "block")
53
        self.model_type = model_type
helloyongyang's avatar
helloyongyang committed
54
55
56
57
58

        if self.config["seq_parallel"]:
            self.seq_p_group = self.config.get("device_mesh").get_group(mesh_dim="seq_p")
        else:
            self.seq_p_group = None
59

gushiqiao's avatar
gushiqiao committed
60
        self.clean_cuda_cache = self.config.get("clean_cuda_cache", False)
61
        self.dit_quantized = self.config.get("dit_quantized", False)
62
        if self.dit_quantized:
63
            assert self.config.get("dit_quant_scheme", "Default") in ["Default-Force-FP32", "fp8-vllm", "int8-vllm", "fp8-q8f", "int8-q8f", "fp8-b128-deepgemm", "fp8-sgl", "int8-sgl", "int8-torchao"]
gushiqiao's avatar
gushiqiao committed
64
        self.device = device
helloyongyang's avatar
helloyongyang committed
65
66
67
68
69
70
71
        self._init_infer_class()
        self._init_weights()
        self._init_infer()

    def _init_infer_class(self):
        self.pre_infer_class = WanPreInfer
        self.post_infer_class = WanPostInfer
helloyongyang's avatar
helloyongyang committed
72
73

        if self.config["feature_caching"] == "NoCaching":
74
            self.transformer_infer_class = WanTransformerInfer if not self.cpu_offload else WanOffloadTransformerInfer
helloyongyang's avatar
helloyongyang committed
75
76
77
78
79
80
81
82
83
84
85
86
87
88
        elif self.config["feature_caching"] == "Tea":
            self.transformer_infer_class = WanTransformerInferTeaCaching
        elif self.config["feature_caching"] == "TaylorSeer":
            self.transformer_infer_class = WanTransformerInferTaylorCaching
        elif self.config["feature_caching"] == "Ada":
            self.transformer_infer_class = WanTransformerInferAdaCaching
        elif self.config["feature_caching"] == "Custom":
            self.transformer_infer_class = WanTransformerInferCustomCaching
        elif self.config["feature_caching"] == "FirstBlock":
            self.transformer_infer_class = WanTransformerInferFirstBlock
        elif self.config["feature_caching"] == "DualBlock":
            self.transformer_infer_class = WanTransformerInferDualBlock
        elif self.config["feature_caching"] == "DynamicBlock":
            self.transformer_infer_class = WanTransformerInferDynamicBlock
Musisoul's avatar
Musisoul committed
89
90
        elif self.config["feature_caching"] == "Mag":
            self.transformer_infer_class = WanTransformerInferMagCaching
helloyongyang's avatar
helloyongyang committed
91
        else:
helloyongyang's avatar
helloyongyang committed
92
            raise NotImplementedError(f"Unsupported feature_caching type: {self.config['feature_caching']}")
helloyongyang's avatar
helloyongyang committed
93

gushiqiao's avatar
gushiqiao committed
94
95
96
97
98
99
    def _should_load_weights(self):
        """Determine if current rank should load weights from disk."""
        if self.config.get("device_mesh") is None:
            # Single GPU mode
            return True
        elif dist.is_initialized():
100
101
102
103
104
105
            if self.config.get("load_from_rank0", False):
                # Multi-GPU mode, only rank 0 loads
                if dist.get_rank() == 0:
                    logger.info(f"Loading weights from {self.model_path}")
                    return True
            else:
gushiqiao's avatar
gushiqiao committed
106
107
108
                return True
        return False

109
    def _should_init_empty_model(self):
110
        if self.config.get("lora_configs") and self.config["lora_configs"]:
111
112
113
114
115
116
117
118
119
120
121
122
            if self.model_type in ["wan2.1"]:
                return True
            if self.model_type in ["wan2.2_moe_high_noise"]:
                for lora_config in self.config["lora_configs"]:
                    if lora_config["name"] == "high_noise_model":
                        return True
            if self.model_type in ["wan2.2_moe_low_noise"]:
                for lora_config in self.config["lora_configs"]:
                    if lora_config["name"] == "low_noise_model":
                        return True
        return False

123
    def _load_safetensor_to_dict(self, file_path, unified_dtype, sensitive_layer):
124
125
        remove_keys = self.remove_keys if hasattr(self, "remove_keys") else []

126
127
128
129
        if self.device.type == "cuda" and dist.is_initialized():
            device = torch.device("cuda:{}".format(dist.get_rank()))
        else:
            device = self.device
130

131
        with safe_open(file_path, framework="pt", device=str(device)) as f:
132
133
134
135
136
            return {
                key: (f.get_tensor(key).to(GET_DTYPE()) if unified_dtype or all(s not in key for s in sensitive_layer) else f.get_tensor(key).to(GET_SENSITIVE_DTYPE()))
                for key in f.keys()
                if not any(remove_key in key for remove_key in remove_keys)
            }
helloyongyang's avatar
helloyongyang committed
137

138
    def _load_ckpt(self, unified_dtype, sensitive_layer):
139
140
141
142
143
144
145
146
147
        if self.config.get("dit_original_ckpt", None):
            safetensors_path = self.config["dit_original_ckpt"]
        else:
            safetensors_path = self.model_path

        if os.path.isdir(safetensors_path):
            safetensors_files = glob.glob(os.path.join(safetensors_path, "*.safetensors"))
        else:
            safetensors_files = [safetensors_path]
148

helloyongyang's avatar
helloyongyang committed
149
150
        weight_dict = {}
        for file_path in safetensors_files:
151
            if self.config.get("adapter_model_path", None) is not None:
152
                if self.config["adapter_model_path"] == file_path:
153
                    continue
154
            logger.info(f"Loading weights from {file_path}")
155
            file_weights = self._load_safetensor_to_dict(file_path, unified_dtype, sensitive_layer)
helloyongyang's avatar
helloyongyang committed
156
            weight_dict.update(file_weights)
157

helloyongyang's avatar
helloyongyang committed
158
159
        return weight_dict

160
    def _load_quant_ckpt(self, unified_dtype, sensitive_layer):
161
        remove_keys = self.remove_keys if hasattr(self, "remove_keys") else []
gushiqiao's avatar
Fix  
gushiqiao committed
162

163
164
165
166
        if self.config.get("dit_quantized_ckpt", None):
            safetensors_path = self.config["dit_quantized_ckpt"]
        else:
            safetensors_path = self.model_path
gushiqiao's avatar
Fix  
gushiqiao committed
167

168
169
170
171
        if os.path.isdir(safetensors_path):
            safetensors_files = glob.glob(os.path.join(safetensors_path, "*.safetensors"))
        else:
            safetensors_files = [safetensors_path]
gushiqiao's avatar
Fix  
gushiqiao committed
172
173

        weight_dict = {}
174
175
176
177
        for safetensor_path in safetensors_files:
            if self.config.get("adapter_model_path", None) is not None:
                if self.config["adapter_model_path"] == safetensor_path:
                    continue
gushiqiao's avatar
Fix  
gushiqiao committed
178
179
180
            with safe_open(safetensor_path, framework="pt") as f:
                logger.info(f"Loading weights from {safetensor_path}")
                for k in f.keys():
181
182
                    if any(remove_key in k for remove_key in remove_keys):
                        continue
183
184
185
186
187
                    if f.get_tensor(k).dtype in [
                        torch.float16,
                        torch.bfloat16,
                        torch.float,
                    ]:
188
                        if unified_dtype or all(s not in k for s in sensitive_layer):
gushiqiao's avatar
gushiqiao committed
189
                            weight_dict[k] = f.get_tensor(k).to(GET_DTYPE()).to(self.device)
gushiqiao's avatar
Fix  
gushiqiao committed
190
                        else:
gushiqiao's avatar
gushiqiao committed
191
                            weight_dict[k] = f.get_tensor(k).to(GET_SENSITIVE_DTYPE()).to(self.device)
gushiqiao's avatar
Fix  
gushiqiao committed
192
                    else:
gushiqiao's avatar
gushiqiao committed
193
                        weight_dict[k] = f.get_tensor(k).to(self.device)
194

195
196
        return weight_dict

197
    def _load_quant_split_ckpt(self, unified_dtype, sensitive_layer):  # Need rewrite
gushiqiao's avatar
gushiqiao committed
198
        lazy_load_model_path = self.dit_quantized_ckpt
199
        logger.info(f"Loading splited quant model from {lazy_load_model_path}")
gushiqiao's avatar
gushiqiao committed
200
        pre_post_weight_dict = {}
201
202

        safetensor_path = os.path.join(lazy_load_model_path, "non_block.safetensors")
gushiqiao's avatar
gushiqiao committed
203
        with safe_open(safetensor_path, framework="pt", device="cpu") as f:
204
            for k in f.keys():
205
206
207
208
209
                if f.get_tensor(k).dtype in [
                    torch.float16,
                    torch.bfloat16,
                    torch.float,
                ]:
210
                    if unified_dtype or all(s not in k for s in sensitive_layer):
gushiqiao's avatar
gushiqiao committed
211
                        pre_post_weight_dict[k] = f.get_tensor(k).to(GET_DTYPE()).to(self.device)
gushiqiao's avatar
Fix  
gushiqiao committed
212
                    else:
gushiqiao's avatar
gushiqiao committed
213
                        pre_post_weight_dict[k] = f.get_tensor(k).to(GET_SENSITIVE_DTYPE()).to(self.device)
gushiqiao's avatar
Fix  
gushiqiao committed
214
                else:
gushiqiao's avatar
gushiqiao committed
215
                    pre_post_weight_dict[k] = f.get_tensor(k).to(self.device)
216

gushiqiao's avatar
gushiqiao committed
217
        return pre_post_weight_dict
218

219
220
221
222
223
224
225
226
    def _load_gguf_ckpt(self):
        gguf_path = self.dit_quantized_ckpt
        logger.info(f"Loading gguf-quant dit model from {gguf_path}")
        reader = gguf.GGUFReader(gguf_path)
        for tensor in reader.tensors:
            # TODO: implement _load_gguf_ckpt
            pass

lijiaqi2's avatar
lijiaqi2 committed
227
    def _init_weights(self, weight_dict=None):
228
        unified_dtype = GET_DTYPE() == GET_SENSITIVE_DTYPE()
gushiqiao's avatar
Fix  
gushiqiao committed
229
        # Some layers run with float32 to achieve high accuracy
230
        sensitive_layer = {
gushiqiao's avatar
gushiqiao committed
231
232
233
234
235
236
            "norm",
            "embedding",
            "modulation",
            "time",
            "img_emb.proj.0",
            "img_emb.proj.4",
gushiqiao's avatar
gushiqiao committed
237
238
            "before_proj",  # vace
            "after_proj",  # vace
gushiqiao's avatar
gushiqiao committed
239
        }
240

lijiaqi2's avatar
lijiaqi2 committed
241
        if weight_dict is None:
gushiqiao's avatar
gushiqiao committed
242
            is_weight_loader = self._should_load_weights()
243
            if is_weight_loader:
244
                if not self.dit_quantized:
gushiqiao's avatar
gushiqiao committed
245
246
                    # Load original weights
                    weight_dict = self._load_ckpt(unified_dtype, sensitive_layer)
247
                else:
gushiqiao's avatar
gushiqiao committed
248
                    # Load quantized weights
249
                    if not self.config.get("lazy_load", False):
gushiqiao's avatar
gushiqiao committed
250
                        weight_dict = self._load_quant_ckpt(unified_dtype, sensitive_layer)
251
                    else:
gushiqiao's avatar
gushiqiao committed
252
                        weight_dict = self._load_quant_split_ckpt(unified_dtype, sensitive_layer)
253

254
255
            if self.config.get("device_mesh") is not None and self.config.get("load_from_rank0", False):
                weight_dict = self._load_weights_from_rank0(weight_dict, is_weight_loader)
256

257
258
259
            if hasattr(self, "adapter_weights_dict"):
                weight_dict.update(self.adapter_weights_dict)

gushiqiao's avatar
gushiqiao committed
260
            self.original_weight_dict = weight_dict
lijiaqi2's avatar
lijiaqi2 committed
261
262
        else:
            self.original_weight_dict = weight_dict
263

gushiqiao's avatar
gushiqiao committed
264
        # Initialize weight containers
helloyongyang's avatar
helloyongyang committed
265
266
        self.pre_weight = self.pre_weight_class(self.config)
        self.transformer_weights = self.transformer_weight_class(self.config)
267
        if not self._should_init_empty_model():
268
            self._apply_weights()
gushiqiao's avatar
gushiqiao committed
269

270
271
272
273
274
    def _apply_weights(self, weight_dict=None):
        if weight_dict is not None:
            self.original_weight_dict = weight_dict
            del weight_dict
            gc.collect()
gushiqiao's avatar
gushiqiao committed
275
        # Load weights into containers
276
        self.pre_weight.load(self.original_weight_dict)
gushiqiao's avatar
gushiqiao committed
277
        self.transformer_weights.load(self.original_weight_dict)
helloyongyang's avatar
helloyongyang committed
278

gushiqiao's avatar
gushiqiao committed
279
280
281
282
        del self.original_weight_dict
        torch.cuda.empty_cache()
        gc.collect()

283
284
    def _load_weights_from_rank0(self, weight_dict, is_weight_loader):
        logger.info("Loading distributed weights")
gushiqiao's avatar
gushiqiao committed
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
        global_src_rank = 0
        target_device = "cpu" if self.cpu_offload else "cuda"

        if is_weight_loader:
            meta_dict = {}
            for key, tensor in weight_dict.items():
                meta_dict[key] = {"shape": tensor.shape, "dtype": tensor.dtype}

            obj_list = [meta_dict]
            dist.broadcast_object_list(obj_list, src=global_src_rank)
            synced_meta_dict = obj_list[0]
        else:
            obj_list = [None]
            dist.broadcast_object_list(obj_list, src=global_src_rank)
            synced_meta_dict = obj_list[0]

        distributed_weight_dict = {}
        for key, meta in synced_meta_dict.items():
            distributed_weight_dict[key] = torch.empty(meta["shape"], dtype=meta["dtype"], device=target_device)

        if target_device == "cuda":
            dist.barrier(device_ids=[torch.cuda.current_device()])

        for key in sorted(synced_meta_dict.keys()):
            if is_weight_loader:
                distributed_weight_dict[key].copy_(weight_dict[key], non_blocking=True)

gushiqiao's avatar
gushiqiao committed
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
            if target_device == "cpu":
                if is_weight_loader:
                    gpu_tensor = distributed_weight_dict[key].cuda()
                    dist.broadcast(gpu_tensor, src=global_src_rank)
                    distributed_weight_dict[key].copy_(gpu_tensor.cpu(), non_blocking=True)
                    del gpu_tensor
                    torch.cuda.empty_cache()
                else:
                    gpu_tensor = torch.empty_like(distributed_weight_dict[key], device="cuda")
                    dist.broadcast(gpu_tensor, src=global_src_rank)
                    distributed_weight_dict[key].copy_(gpu_tensor.cpu(), non_blocking=True)
                    del gpu_tensor
                    torch.cuda.empty_cache()

                if distributed_weight_dict[key].is_pinned():
                    distributed_weight_dict[key].copy_(distributed_weight_dict[key], non_blocking=True)
            else:
                dist.broadcast(distributed_weight_dict[key], src=global_src_rank)

        if target_device == "cuda":
            torch.cuda.synchronize()
        else:
            for tensor in distributed_weight_dict.values():
                if tensor.is_pinned():
                    tensor.copy_(tensor, non_blocking=False)
gushiqiao's avatar
gushiqiao committed
337
338

        logger.info(f"Weights distributed across {dist.get_world_size()} devices on {target_device}")
339

gushiqiao's avatar
gushiqiao committed
340
341
        return distributed_weight_dict

helloyongyang's avatar
helloyongyang committed
342
343
344
    def _init_infer(self):
        self.pre_infer = self.pre_infer_class(self.config)
        self.post_infer = self.post_infer_class(self.config)
helloyongyang's avatar
helloyongyang committed
345
        self.transformer_infer = self.transformer_infer_class(self.config)
helloyongyang's avatar
helloyongyang committed
346
347
348

    def set_scheduler(self, scheduler):
        self.scheduler = scheduler
349
350
        self.pre_infer.set_scheduler(scheduler)
        self.post_infer.set_scheduler(scheduler)
helloyongyang's avatar
helloyongyang committed
351
352
        self.transformer_infer.set_scheduler(scheduler)

TorynCurtis's avatar
TorynCurtis committed
353
354
355
356
357
358
359
360
    def to_cpu(self):
        self.pre_weight.to_cpu()
        self.transformer_weights.to_cpu()

    def to_cuda(self):
        self.pre_weight.to_cuda()
        self.transformer_weights.to_cuda()

helloyongyang's avatar
helloyongyang committed
361
362
    @torch.no_grad()
    def infer(self, inputs):
363
        if self.cpu_offload:
364
            if self.offload_granularity == "model" and self.scheduler.step_index == 0 and "wan2.2_moe" not in self.config["model_cls"]:
365
366
367
                self.to_cuda()
            elif self.offload_granularity != "model":
                self.pre_weight.to_cuda()
gushiqiao's avatar
gushiqiao committed
368
                self.transformer_weights.non_block_weights_to_cuda()
369

370
        if self.config["enable_cfg"]:
helloyongyang's avatar
helloyongyang committed
371
372
373
374
375
376
377
            if self.config["cfg_parallel"]:
                # ==================== CFG Parallel Processing ====================
                cfg_p_group = self.config["device_mesh"].get_group(mesh_dim="cfg_p")
                assert dist.get_world_size(cfg_p_group) == 2, "cfg_p_world_size must be equal to 2"
                cfg_p_rank = dist.get_rank(cfg_p_group)

                if cfg_p_rank == 0:
helloyongyang's avatar
helloyongyang committed
378
                    noise_pred = self._infer_cond_uncond(inputs, infer_condition=True)
helloyongyang's avatar
helloyongyang committed
379
                else:
helloyongyang's avatar
helloyongyang committed
380
                    noise_pred = self._infer_cond_uncond(inputs, infer_condition=False)
helloyongyang's avatar
helloyongyang committed
381

helloyongyang's avatar
helloyongyang committed
382
383
384
385
386
387
                noise_pred_list = [torch.zeros_like(noise_pred) for _ in range(2)]
                dist.all_gather(noise_pred_list, noise_pred, group=cfg_p_group)
                noise_pred_cond = noise_pred_list[0]  # cfg_p_rank == 0
                noise_pred_uncond = noise_pred_list[1]  # cfg_p_rank == 1
            else:
                # ==================== CFG Processing ====================
helloyongyang's avatar
helloyongyang committed
388
389
                noise_pred_cond = self._infer_cond_uncond(inputs, infer_condition=True)
                noise_pred_uncond = self._infer_cond_uncond(inputs, infer_condition=False)
gushiqiao's avatar
gushiqiao committed
390

helloyongyang's avatar
helloyongyang committed
391
392
393
            self.scheduler.noise_pred = noise_pred_uncond + self.scheduler.sample_guide_scale * (noise_pred_cond - noise_pred_uncond)
        else:
            # ==================== No CFG ====================
helloyongyang's avatar
helloyongyang committed
394
            self.scheduler.noise_pred = self._infer_cond_uncond(inputs, infer_condition=True)
395
396

        if self.cpu_offload:
397
            if self.offload_granularity == "model" and self.scheduler.step_index == self.scheduler.infer_steps - 1 and "wan2.2_moe" not in self.config["model_cls"]:
398
399
                self.to_cpu()
            elif self.offload_granularity != "model":
root's avatar
root committed
400
                self.pre_weight.to_cpu()
gushiqiao's avatar
gushiqiao committed
401
                self.transformer_weights.non_block_weights_to_cpu()
gushiqiao's avatar
gushiqiao committed
402

Yang Yong(雍洋)'s avatar
Yang Yong(雍洋) committed
403
    @compiled_method()
404
    @torch.no_grad()
helloyongyang's avatar
helloyongyang committed
405
406
407
408
    def _infer_cond_uncond(self, inputs, infer_condition=True):
        self.scheduler.infer_condition = infer_condition

        pre_infer_out = self.pre_infer.infer(self.pre_weight, inputs)
helloyongyang's avatar
helloyongyang committed
409
410
411
412
413
414
415
416
417

        if self.config["seq_parallel"]:
            pre_infer_out = self._seq_parallel_pre_process(pre_infer_out)

        x = self.transformer_infer.infer(self.transformer_weights, pre_infer_out)

        if self.config["seq_parallel"]:
            x = self._seq_parallel_post_process(x)

gushiqiao's avatar
gushiqiao committed
418
        noise_pred = self.post_infer.infer(x, pre_infer_out)[0]
helloyongyang's avatar
helloyongyang committed
419
420
421
422
423
424
425
426
427

        if self.clean_cuda_cache:
            del x, pre_infer_out
            torch.cuda.empty_cache()

        return noise_pred

    @torch.no_grad()
    def _seq_parallel_pre_process(self, pre_infer_out):
helloyongyang's avatar
helloyongyang committed
428
        x = pre_infer_out.x
helloyongyang's avatar
helloyongyang committed
429
430
431
432
433
        world_size = dist.get_world_size(self.seq_p_group)
        cur_rank = dist.get_rank(self.seq_p_group)

        padding_size = (world_size - (x.shape[0] % world_size)) % world_size
        if padding_size > 0:
helloyongyang's avatar
helloyongyang committed
434
            x = F.pad(x, (0, 0, 0, padding_size))
helloyongyang's avatar
helloyongyang committed
435

helloyongyang's avatar
helloyongyang committed
436
        pre_infer_out.x = torch.chunk(x, world_size, dim=0)[cur_rank]
helloyongyang's avatar
helloyongyang committed
437

438
        if self.config["model_cls"] in ["wan2.2", "wan2.2_audio"] and self.config["task"] in ["i2v", "s2v"]:
helloyongyang's avatar
helloyongyang committed
439
440
441
442
443
444
445
            embed, embed0 = pre_infer_out.embed, pre_infer_out.embed0

            padding_size = (world_size - (embed.shape[0] % world_size)) % world_size
            if padding_size > 0:
                embed = F.pad(embed, (0, 0, 0, padding_size))
                embed0 = F.pad(embed0, (0, 0, 0, 0, 0, padding_size))

helloyongyang's avatar
helloyongyang committed
446
447
            pre_infer_out.embed = torch.chunk(embed, world_size, dim=0)[cur_rank]
            pre_infer_out.embed0 = torch.chunk(embed0, world_size, dim=0)[cur_rank]
helloyongyang's avatar
helloyongyang committed
448
449
450
451
452
453
454
455
456

        return pre_infer_out

    @torch.no_grad()
    def _seq_parallel_post_process(self, x):
        world_size = dist.get_world_size(self.seq_p_group)
        gathered_x = [torch.empty_like(x) for _ in range(world_size)]
        dist.all_gather(gathered_x, x, group=self.seq_p_group)
        combined_output = torch.cat(gathered_x, dim=0)
helloyongyang's avatar
helloyongyang committed
457
        return combined_output