model.py 18.7 KB
Newer Older
gushiqiao's avatar
gushiqiao committed
1
import gc
2
3
import os

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

10
from lightx2v.models.networks.wan.infer.feature_caching.transformer_infer import (
11
12
    WanTransformerInferAdaCaching,
    WanTransformerInferCustomCaching,
Rongjin Yang's avatar
Rongjin Yang committed
13
14
    WanTransformerInferDualBlock,
    WanTransformerInferDynamicBlock,
PengGao's avatar
PengGao committed
15
    WanTransformerInferFirstBlock,
Musisoul's avatar
Musisoul committed
16
    WanTransformerInferMagCaching,
PengGao's avatar
PengGao committed
17
18
19
    WanTransformerInferTaylorCaching,
    WanTransformerInferTeaCaching,
)
20
21
22
from lightx2v.models.networks.wan.infer.offload.transformer_infer import (
    WanOffloadTransformerInfer,
)
PengGao's avatar
PengGao committed
23
24
25
26
27
28
29
30
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,
31
)
Yang Yong(雍洋)'s avatar
Yang Yong(雍洋) committed
32
from lightx2v.utils.custom_compiler import CompiledMethodsMixin, compiled_method
33
from lightx2v.utils.envs import *
34
from lightx2v.utils.utils import *
helloyongyang's avatar
helloyongyang committed
35

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

helloyongyang's avatar
helloyongyang committed
41

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

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

        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
57

58
59
60
61
62
        if self.config.get("lora_configs") and self.config.lora_configs:
            self.init_empty_model = True
        else:
            self.init_empty_model = False

gushiqiao's avatar
gushiqiao committed
63
        self.clean_cuda_cache = self.config.get("clean_cuda_cache", False)
64
        self.dit_quantized = self.config.get("dit_quantized", False)
65
        if self.dit_quantized:
66
            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
67
        self.device = device
helloyongyang's avatar
helloyongyang committed
68
69
70
71
72
73
74
        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
75
76

        if self.config["feature_caching"] == "NoCaching":
77
            self.transformer_infer_class = WanTransformerInfer if not self.cpu_offload else WanOffloadTransformerInfer
helloyongyang's avatar
helloyongyang committed
78
79
80
81
82
83
84
85
86
87
88
89
90
91
        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
92
93
        elif self.config["feature_caching"] == "Mag":
            self.transformer_infer_class = WanTransformerInferMagCaching
helloyongyang's avatar
helloyongyang committed
94
        else:
helloyongyang's avatar
helloyongyang committed
95
            raise NotImplementedError(f"Unsupported feature_caching type: {self.config['feature_caching']}")
helloyongyang's avatar
helloyongyang committed
96

gushiqiao's avatar
gushiqiao committed
97
98
99
100
101
102
    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():
103
104
105
106
107
108
            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
109
110
111
                return True
        return False

112
    def _load_safetensor_to_dict(self, file_path, unified_dtype, sensitive_layer):
113
114
        remove_keys = self.remove_keys if hasattr(self, "remove_keys") else []

115
116
117
118
        if self.device.type == "cuda" and dist.is_initialized():
            device = torch.device("cuda:{}".format(dist.get_rank()))
        else:
            device = self.device
119

120
        with safe_open(file_path, framework="pt", device=str(device)) as f:
121
122
123
124
125
            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
126

127
    def _load_ckpt(self, unified_dtype, sensitive_layer):
128
129
130
131
132
133
134
135
136
        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]
137

helloyongyang's avatar
helloyongyang committed
138
139
        weight_dict = {}
        for file_path in safetensors_files:
140
            if self.config.get("adapter_model_path", None) is not None:
141
                if self.config["adapter_model_path"] == file_path:
142
                    continue
143
            logger.info(f"Loading weights from {file_path}")
144
            file_weights = self._load_safetensor_to_dict(file_path, unified_dtype, sensitive_layer)
helloyongyang's avatar
helloyongyang committed
145
            weight_dict.update(file_weights)
146

helloyongyang's avatar
helloyongyang committed
147
148
        return weight_dict

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

152
153
154
155
        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
156

157
158
159
160
        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
161
162

        weight_dict = {}
163
164
165
166
        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
167
168
169
            with safe_open(safetensor_path, framework="pt") as f:
                logger.info(f"Loading weights from {safetensor_path}")
                for k in f.keys():
170
171
                    if any(remove_key in k for remove_key in remove_keys):
                        continue
172
173
174
175
176
                    if f.get_tensor(k).dtype in [
                        torch.float16,
                        torch.bfloat16,
                        torch.float,
                    ]:
177
                        if unified_dtype or all(s not in k for s in sensitive_layer):
gushiqiao's avatar
gushiqiao committed
178
                            weight_dict[k] = f.get_tensor(k).to(GET_DTYPE()).to(self.device)
gushiqiao's avatar
Fix  
gushiqiao committed
179
                        else:
gushiqiao's avatar
gushiqiao committed
180
                            weight_dict[k] = f.get_tensor(k).to(GET_SENSITIVE_DTYPE()).to(self.device)
gushiqiao's avatar
Fix  
gushiqiao committed
181
                    else:
gushiqiao's avatar
gushiqiao committed
182
                        weight_dict[k] = f.get_tensor(k).to(self.device)
183

184
185
        return weight_dict

186
    def _load_quant_split_ckpt(self, unified_dtype, sensitive_layer):  # Need rewrite
gushiqiao's avatar
gushiqiao committed
187
        lazy_load_model_path = self.dit_quantized_ckpt
188
        logger.info(f"Loading splited quant model from {lazy_load_model_path}")
gushiqiao's avatar
gushiqiao committed
189
        pre_post_weight_dict = {}
190
191

        safetensor_path = os.path.join(lazy_load_model_path, "non_block.safetensors")
gushiqiao's avatar
gushiqiao committed
192
        with safe_open(safetensor_path, framework="pt", device="cpu") as f:
193
            for k in f.keys():
194
195
196
197
198
                if f.get_tensor(k).dtype in [
                    torch.float16,
                    torch.bfloat16,
                    torch.float,
                ]:
199
                    if unified_dtype or all(s not in k for s in sensitive_layer):
gushiqiao's avatar
gushiqiao committed
200
                        pre_post_weight_dict[k] = f.get_tensor(k).to(GET_DTYPE()).to(self.device)
gushiqiao's avatar
Fix  
gushiqiao committed
201
                    else:
gushiqiao's avatar
gushiqiao committed
202
                        pre_post_weight_dict[k] = f.get_tensor(k).to(GET_SENSITIVE_DTYPE()).to(self.device)
gushiqiao's avatar
Fix  
gushiqiao committed
203
                else:
gushiqiao's avatar
gushiqiao committed
204
                    pre_post_weight_dict[k] = f.get_tensor(k).to(self.device)
205

gushiqiao's avatar
gushiqiao committed
206
        return pre_post_weight_dict
207

208
209
210
211
212
213
214
215
    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
216
    def _init_weights(self, weight_dict=None):
217
        unified_dtype = GET_DTYPE() == GET_SENSITIVE_DTYPE()
gushiqiao's avatar
Fix  
gushiqiao committed
218
        # Some layers run with float32 to achieve high accuracy
219
        sensitive_layer = {
gushiqiao's avatar
gushiqiao committed
220
221
222
223
224
225
            "norm",
            "embedding",
            "modulation",
            "time",
            "img_emb.proj.0",
            "img_emb.proj.4",
gushiqiao's avatar
gushiqiao committed
226
227
            "before_proj",  # vace
            "after_proj",  # vace
gushiqiao's avatar
gushiqiao committed
228
        }
229

lijiaqi2's avatar
lijiaqi2 committed
230
        if weight_dict is None:
gushiqiao's avatar
gushiqiao committed
231
            is_weight_loader = self._should_load_weights()
232
            if is_weight_loader:
233
                if not self.dit_quantized:
gushiqiao's avatar
gushiqiao committed
234
235
                    # Load original weights
                    weight_dict = self._load_ckpt(unified_dtype, sensitive_layer)
236
                else:
gushiqiao's avatar
gushiqiao committed
237
                    # Load quantized weights
238
                    if not self.config.get("lazy_load", False):
gushiqiao's avatar
gushiqiao committed
239
                        weight_dict = self._load_quant_ckpt(unified_dtype, sensitive_layer)
240
                    else:
gushiqiao's avatar
gushiqiao committed
241
                        weight_dict = self._load_quant_split_ckpt(unified_dtype, sensitive_layer)
242

243
244
            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)
245

246
247
248
            if hasattr(self, "adapter_weights_dict"):
                weight_dict.update(self.adapter_weights_dict)

gushiqiao's avatar
gushiqiao committed
249
            self.original_weight_dict = weight_dict
lijiaqi2's avatar
lijiaqi2 committed
250
251
        else:
            self.original_weight_dict = weight_dict
252

gushiqiao's avatar
gushiqiao committed
253
        # Initialize weight containers
helloyongyang's avatar
helloyongyang committed
254
255
        self.pre_weight = self.pre_weight_class(self.config)
        self.transformer_weights = self.transformer_weight_class(self.config)
256
257
        if not self.init_empty_model:
            self._apply_weights()
gushiqiao's avatar
gushiqiao committed
258

259
260
261
262
263
    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
264
        # Load weights into containers
265
        self.pre_weight.load(self.original_weight_dict)
gushiqiao's avatar
gushiqiao committed
266
        self.transformer_weights.load(self.original_weight_dict)
helloyongyang's avatar
helloyongyang committed
267

gushiqiao's avatar
gushiqiao committed
268
269
270
271
        del self.original_weight_dict
        torch.cuda.empty_cache()
        gc.collect()

272
273
    def _load_weights_from_rank0(self, weight_dict, is_weight_loader):
        logger.info("Loading distributed weights")
gushiqiao's avatar
gushiqiao committed
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
        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
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
            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
326
327

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

gushiqiao's avatar
gushiqiao committed
329
330
        return distributed_weight_dict

helloyongyang's avatar
helloyongyang committed
331
332
333
    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
334
        self.transformer_infer = self.transformer_infer_class(self.config)
helloyongyang's avatar
helloyongyang committed
335
336
337

    def set_scheduler(self, scheduler):
        self.scheduler = scheduler
338
339
        self.pre_infer.set_scheduler(scheduler)
        self.post_infer.set_scheduler(scheduler)
helloyongyang's avatar
helloyongyang committed
340
341
        self.transformer_infer.set_scheduler(scheduler)

TorynCurtis's avatar
TorynCurtis committed
342
343
344
345
346
347
348
349
    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
350
351
    @torch.no_grad()
    def infer(self, inputs):
352
        if self.cpu_offload:
353
            if self.offload_granularity == "model" and self.scheduler.step_index == 0 and "wan2.2_moe" not in self.config["model_cls"]:
354
355
356
                self.to_cuda()
            elif self.offload_granularity != "model":
                self.pre_weight.to_cuda()
gushiqiao's avatar
gushiqiao committed
357
                self.transformer_weights.non_block_weights_to_cuda()
358

359
        if self.config["enable_cfg"]:
helloyongyang's avatar
helloyongyang committed
360
361
362
363
364
365
366
            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
367
                    noise_pred = self._infer_cond_uncond(inputs, infer_condition=True)
helloyongyang's avatar
helloyongyang committed
368
                else:
helloyongyang's avatar
helloyongyang committed
369
                    noise_pred = self._infer_cond_uncond(inputs, infer_condition=False)
helloyongyang's avatar
helloyongyang committed
370

helloyongyang's avatar
helloyongyang committed
371
372
373
374
375
376
                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
377
378
                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
379

helloyongyang's avatar
helloyongyang committed
380
381
382
            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
383
            self.scheduler.noise_pred = self._infer_cond_uncond(inputs, infer_condition=True)
384
385

        if self.cpu_offload:
386
            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"]:
387
388
                self.to_cpu()
            elif self.offload_granularity != "model":
root's avatar
root committed
389
                self.pre_weight.to_cpu()
gushiqiao's avatar
gushiqiao committed
390
                self.transformer_weights.non_block_weights_to_cpu()
gushiqiao's avatar
gushiqiao committed
391

Yang Yong(雍洋)'s avatar
Yang Yong(雍洋) committed
392
    @compiled_method()
393
    @torch.no_grad()
helloyongyang's avatar
helloyongyang committed
394
395
396
397
    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
398
399
400
401
402
403
404
405
406

        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
407
        noise_pred = self.post_infer.infer(x, pre_infer_out)[0]
helloyongyang's avatar
helloyongyang committed
408
409
410
411
412
413
414
415
416

        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
417
        x = pre_infer_out.x
helloyongyang's avatar
helloyongyang committed
418
419
420
421
422
        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
423
            x = F.pad(x, (0, 0, 0, padding_size))
helloyongyang's avatar
helloyongyang committed
424

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

427
        if self.config["model_cls"] in ["wan2.2", "wan2.2_audio"] and self.config["task"] in ["i2v", "s2v"]:
helloyongyang's avatar
helloyongyang committed
428
429
430
431
432
433
434
            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
435
436
            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
437
438
439
440
441
442
443
444
445

        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
446
        return combined_output