+
+
+
+
+
+
@@ -310,4 +310,14 @@ If you wish to cite relevant research papars, you can find the reference below.
journal={arXiv},
year={2023}
}
+
+# Distrifusion
+@InProceedings{Li_2024_CVPR,
+ author={Li, Muyang and Cai, Tianle and Cao, Jiaxin and Zhang, Qinsheng and Cai, Han and Bai, Junjie and Jia, Yangqing and Li, Kai and Han, Song},
+ title={DistriFusion: Distributed Parallel Inference for High-Resolution Diffusion Models},
+ booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
+ month={June},
+ year={2024},
+ pages={7183-7193}
+}
```
diff --git a/colossalai/inference/config.py b/colossalai/inference/config.py
index e114e8a61..072ddbcfd 100644
--- a/colossalai/inference/config.py
+++ b/colossalai/inference/config.py
@@ -5,7 +5,7 @@ Our config contains various options for inference optimization, it is a unified
import logging
from abc import ABC, abstractmethod
from dataclasses import dataclass, fields
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Union
import torch
from transformers.generation import GenerationConfig
@@ -186,6 +186,7 @@ class InferenceConfig(RPC_PARAM):
enable_streamingllm(bool): Whether to use StreamingLLM, the relevant algorithms refer to the paper at https://arxiv.org/pdf/2309.17453 for implementation.
start_token_size(int): The size of the start tokens, when using StreamingLLM.
generated_token_size(int): The size of the generated tokens, When using StreamingLLM.
+ patched_parallelism_size(int): Patched Parallelism Size, When using Distrifusion
"""
# NOTE: arrange configs according to their importance and frequency of usage
@@ -245,6 +246,11 @@ class InferenceConfig(RPC_PARAM):
start_token_size: int = 4
generated_token_size: int = 512
+ # Acceleration for Diffusion Model(PipeFusion or Distrifusion)
+ patched_parallelism_size: int = 1 # for distrifusion
+ # pipeFusion_m_size: int = 1 # for pipefusion
+ # pipeFusion_n_size: int = 1 # for pipefusion
+
def __post_init__(self):
self.max_context_len_to_capture = self.max_input_len + self.max_output_len
self._verify_config()
@@ -288,6 +294,14 @@ class InferenceConfig(RPC_PARAM):
# Thereafter, we swap out tokens in units of blocks, and always swapping out the second block when the generated tokens exceeded the limit.
self.start_token_size = self.block_size
+ # check Distrifusion
+ # TODO(@lry89757) need more detailed check
+ if self.patched_parallelism_size > 1:
+ # self.use_patched_parallelism = True
+ self.tp_size = (
+ self.patched_parallelism_size
+ ) # this is not a real tp, because some annoying check, so we have to set this to patched_parallelism_size
+
# check prompt template
if self.prompt_template is None:
return
@@ -324,6 +338,7 @@ class InferenceConfig(RPC_PARAM):
use_cuda_kernel=self.use_cuda_kernel,
use_spec_dec=self.use_spec_dec,
use_flash_attn=use_flash_attn,
+ patched_parallelism_size=self.patched_parallelism_size,
)
return model_inference_config
@@ -396,3 +411,50 @@ class ModelShardInferenceConfig:
use_cuda_kernel: bool = False
use_spec_dec: bool = False
use_flash_attn: bool = False
+ patched_parallelism_size: int = 1 # for diffusion model, Distrifusion Technique
+
+
+@dataclass
+class DiffusionGenerationConfig:
+ """
+ Param for diffusion model forward
+ """
+
+ prompt_2: Optional[Union[str, List[str]]] = None
+ prompt_3: Optional[Union[str, List[str]]] = None
+ height: Optional[int] = None
+ width: Optional[int] = None
+ num_inference_steps: int = None
+ timesteps: List[int] = None
+ guidance_scale: float = None
+ negative_prompt: Optional[Union[str, List[str]]] = (
+ None # NOTE(@lry89757) in pixart default to "", in sd3 default to None
+ )
+ negative_prompt_2: Optional[Union[str, List[str]]] = None
+ negative_prompt_3: Optional[Union[str, List[str]]] = None
+ num_images_per_prompt: Optional[int] = None
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None
+ latents: Optional[torch.FloatTensor] = None
+ prompt_embeds: Optional[torch.FloatTensor] = None
+ negative_prompt_embeds: Optional[torch.FloatTensor] = None
+ pooled_prompt_embeds: Optional[torch.FloatTensor] = None
+ negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None
+ output_type: Optional[str] = None # "pil"
+ return_dict: bool = None
+ joint_attention_kwargs: Optional[Dict[str, Any]] = None
+ clip_skip: Optional[int] = None
+ callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None
+ callback_on_step_end_tensor_inputs: List[str] = None
+
+ def to_dict(self) -> Dict[str, Any]:
+ # NOTE(@lry89757) Only return the dict that not the default value None
+ result = {}
+ for field in fields(self):
+ value = getattr(self, field.name)
+ if value is not None:
+ result[field.name] = value
+ return result
+
+ @classmethod
+ def from_kwargs(cls, **kwargs) -> "DiffusionGenerationConfig":
+ return cls(**kwargs)
diff --git a/colossalai/inference/core/base_engine.py b/colossalai/inference/core/base_engine.py
new file mode 100644
index 000000000..392dd2990
--- /dev/null
+++ b/colossalai/inference/core/base_engine.py
@@ -0,0 +1,90 @@
+from abc import ABC, abstractmethod
+
+import torch
+import torch.nn as nn
+
+from colossalai.cluster import ProcessGroupMesh
+from colossalai.inference.config import ModelShardInferenceConfig
+from colossalai.pipeline.stage_manager import PipelineStageManager
+from colossalai.shardformer import ShardConfig, ShardFormer
+from colossalai.shardformer.policies.base_policy import Policy
+
+
+class BaseEngine(ABC):
+ @abstractmethod
+ def __init__(self, model_or_path, inference_config=None, verbose=False, model_policy=None):
+ pass
+
+ @abstractmethod
+ def init_model(self, model_or_path, model_policy=None, model_shard_infer_config=None):
+ """
+ Init Model for Engine
+ """
+
+ @abstractmethod
+ def generate(self, request_ids=None, prompts=None, generation_config=None, **kwargs):
+ """
+ Generate ouptput for coming requests
+ """
+
+ @abstractmethod
+ def add_request(self, prompts, request_ids=None, **kwargs):
+ """
+ Add new request to Engine
+ """
+
+ @abstractmethod
+ def step(self):
+ """
+ Perform one new step forward
+ """
+
+ @abstractmethod
+ def _verify_args(self):
+ """
+ Verify the parameters and members of class
+ """
+
+ @torch.inference_mode()
+ def capture_model(self):
+ """
+ Use cuda graph to capture model
+ """
+ return NotImplementedError("This method should be implemented by subclasses")
+
+ def _shardformer(
+ self,
+ model: nn.Module,
+ model_policy: Policy,
+ model_shard_infer_config: ModelShardInferenceConfig = None,
+ stage_manager: PipelineStageManager = None,
+ tp_group: ProcessGroupMesh = None,
+ **kwargs,
+ ) -> nn.Module:
+ """
+ Initialize ShardConfig and replace the model with shardformer.
+
+ Args:
+ model (nn.Module): Path or nn.Module of this model.
+ model_policy (Policy): The policy to shardformer model which is determined by the model type.
+ stage_manager (PipelineStageManager, optional): Used to manage pipeline stages. Defaults to None.
+ tp_group (ProcessGroupMesh, optional): Used to manage the process TP group mesh. Defaults to None.
+
+ Returns:
+ nn.Module: The model optimized by Shardformer.
+ """
+
+ shardconfig = ShardConfig(
+ tensor_parallel_process_group=tp_group,
+ pipeline_stage_manager=stage_manager,
+ enable_tensor_parallelism=(self.inference_config.tp_size > 1),
+ enable_fused_normalization=False,
+ enable_all_optimization=False,
+ enable_flash_attention=False,
+ enable_jit_fused=False,
+ enable_sequence_parallelism=False,
+ extra_kwargs={"model_shard_infer_config": model_shard_infer_config, **kwargs},
+ )
+ shardformer = ShardFormer(shard_config=shardconfig)
+ shard_model, _ = shardformer.optimize(model, model_policy)
+ return shard_model
diff --git a/colossalai/inference/core/diffusion_engine.py b/colossalai/inference/core/diffusion_engine.py
new file mode 100644
index 000000000..8bed508cb
--- /dev/null
+++ b/colossalai/inference/core/diffusion_engine.py
@@ -0,0 +1,200 @@
+from itertools import count
+from typing import List, Tuple, Type, Union
+
+import numpy as np
+import PIL.Image
+import torch
+import torch.nn as nn
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from torch import distributed as dist
+
+from colossalai.accelerator import get_accelerator
+from colossalai.cluster import ProcessGroupMesh
+from colossalai.inference.config import DiffusionGenerationConfig, InferenceConfig, ModelShardInferenceConfig
+from colossalai.inference.modeling.layers.diffusion import DiffusionPipe
+from colossalai.inference.modeling.policy import model_policy_map
+from colossalai.inference.struct import DiffusionSequence
+from colossalai.inference.utils import get_model_size, get_model_type
+from colossalai.logging import get_dist_logger
+from colossalai.shardformer.policies.base_policy import Policy
+
+from .base_engine import BaseEngine
+from .request_handler import NaiveRequestHandler
+
+PP_AXIS, TP_AXIS = 0, 1
+
+
+class DiffusionEngine(BaseEngine):
+ def __init__(
+ self,
+ model_or_path: DiffusionPipeline | str,
+ inference_config: InferenceConfig = None,
+ verbose: bool = False,
+ model_policy: Policy | type[Policy] = None,
+ ) -> None:
+ self.inference_config = inference_config
+ self.dtype = inference_config.dtype
+ self.high_precision = inference_config.high_precision
+
+ self.verbose = verbose
+ self.logger = get_dist_logger(__name__)
+ self.model_shard_infer_config = inference_config.to_model_shard_inference_config()
+
+ self.model_type = get_model_type(model_or_path=model_or_path)
+
+ self.init_model(model_or_path, model_policy, self.model_shard_infer_config)
+
+ self.request_handler = NaiveRequestHandler()
+
+ self.counter = count()
+
+ self._verify_args()
+
+ def _verify_args(self) -> None:
+ assert isinstance(self.model, DiffusionPipe), "model must be DiffusionPipe"
+
+ def init_model(
+ self,
+ model_or_path: Union[str, nn.Module, DiffusionPipeline],
+ model_policy: Union[Policy, Type[Policy]] = None,
+ model_shard_infer_config: ModelShardInferenceConfig = None,
+ ):
+ """
+ Shard model or/and Load weight
+
+ Args:
+ model_or_path Union[nn.Module, str]: path to the checkpoint or model of transformer format.
+ model_policy (Policy): the policy to replace the model.
+ model_inference_config: the configuration for modeling initialization when inference.
+ model_shard_infer_config (ModelShardInferenceConfig): the configuration for init of module when inference.
+ """
+ if isinstance(model_or_path, str):
+ model = DiffusionPipeline.from_pretrained(model_or_path, torch_dtype=self.dtype)
+ policy_map_key = model.__class__.__name__
+ model = DiffusionPipe(model)
+ elif isinstance(model_or_path, DiffusionPipeline):
+ policy_map_key = model_or_path.__class__.__name__
+ model = DiffusionPipe(model_or_path)
+ else:
+ self.logger.error(f"model_or_path support only str or DiffusionPipeline currently!")
+
+ torch.cuda.empty_cache()
+ init_gpu_memory = torch.cuda.mem_get_info()[0]
+
+ self.device = get_accelerator().get_current_device()
+ if self.verbose:
+ self.logger.info(f"the device is {self.device}")
+
+ if self.verbose:
+ self.logger.info(
+ f"Before the shard, Rank: [{dist.get_rank()}], model size: {get_model_size(model)} GB, model's device is: {model.device}"
+ )
+
+ if model_policy is None:
+ model_policy = model_policy_map.get(policy_map_key)
+
+ if not isinstance(model_policy, Policy):
+ try:
+ model_policy = model_policy()
+ except Exception as e:
+ raise ValueError(f"Unable to instantiate model policy: {e}")
+
+ assert isinstance(model_policy, Policy), f"Invalid type of model policy: {type(model_policy)}"
+ pg_mesh = ProcessGroupMesh(self.inference_config.pp_size, self.inference_config.tp_size)
+ tp_group = pg_mesh.get_group_along_axis(TP_AXIS)
+
+ self.model = self._shardformer(
+ model,
+ model_policy,
+ model_shard_infer_config,
+ None,
+ tp_group=tp_group,
+ )
+
+ self.model = model.to(self.device)
+
+ if self.verbose:
+ self.logger.info(
+ f"After the shard, Rank: [{dist.get_rank()}], model size: {get_model_size(self.model)} GB, model's device is: {model.device}"
+ )
+
+ free_gpu_memory, _ = torch.cuda.mem_get_info()
+ peak_memory = init_gpu_memory - free_gpu_memory
+ if self.verbose:
+ self.logger.info(
+ f"Rank [{dist.get_rank()}], Model Weight Max Occupy {peak_memory / (1024 ** 3)} GB, Model size: {get_model_size(self.model)} GB"
+ )
+
+ def generate(
+ self,
+ request_ids: Union[List[int], int] = None,
+ prompts: Union[List[str], str] = None,
+ generation_config: DiffusionGenerationConfig = None,
+ **kwargs,
+ ) -> Union[List[Union[str, List[PIL.Image.Image], np.ndarray]], Tuple[List[str], List[List[int]]]]:
+ """ """
+ gen_config_dict = generation_config.to_dict() if generation_config is not None else {}
+ prompts = [prompts] if isinstance(prompts, str) else prompts
+ request_ids = [request_ids] if isinstance(request_ids, int) else request_ids
+
+ with torch.inference_mode():
+ if prompts is not None:
+ self.add_request(
+ request_ids=request_ids,
+ prompts=prompts,
+ **gen_config_dict,
+ **kwargs,
+ )
+
+ output_reqs_list = []
+
+ # intuition: If user provide a generation config, we should replace the existing one.
+ if generation_config is not None:
+ self.generation_config = generation_config
+ self.generation_config_dict = gen_config_dict
+
+ while self.request_handler.check_unfinished_reqs():
+ output_reqs_list += self.step()
+
+ return output_reqs_list
+
+ def add_request(
+ self,
+ prompts: Union[List[str], str],
+ request_ids: Union[List[int], int] = None,
+ **kwargs,
+ ):
+ if request_ids is not None and not isinstance(request_ids, list):
+ request_ids = [request_ids]
+
+ if not isinstance(prompts, list):
+ prompts = [prompts]
+
+ generation_config = DiffusionGenerationConfig.from_kwargs(**kwargs)
+ prompts_num = len(prompts)
+ for i in range(prompts_num):
+ if request_ids:
+ assert isinstance(
+ request_ids[0], int
+ ), f"The request_id type must be int, but got {type(request_ids[0])}"
+ assert len(request_ids) == prompts_num
+ request_id = request_ids[i]
+ else:
+ request_id = next(self.counter)
+
+ seq = DiffusionSequence(request_id=request_id, prompt=prompts[i], generation_config=generation_config)
+
+ self.request_handler.add_sequence(seq)
+
+ def step(self) -> List[PIL.Image.Image]:
+ """
+ In each step, do the follows:
+ 1. Run RequestHandler.schedule() and get the batch used for inference.
+ 2. run forward to get List[Image]
+ Returns:
+ List[PIL.Image.Image]: Image Generated by one step.
+ """
+
+ input = self.request_handler.schedule()
+ ret = self.model(prompt=input.prompt, **input.generation_config.to_dict())
+ return ret
diff --git a/colossalai/inference/core/engine.py b/colossalai/inference/core/engine.py
index 8f8aef65e..5c9bdc321 100644
--- a/colossalai/inference/core/engine.py
+++ b/colossalai/inference/core/engine.py
@@ -1,57 +1,24 @@
-import time
-from itertools import count
-from typing import Dict, List, Optional, Tuple, Type, Union
+from typing import List, Tuple, Type, Union
import numpy as np
-import torch
+import PIL.Image
import torch.nn as nn
-from torch import distributed as dist
-from transformers import (
- AutoConfig,
- AutoModelForCausalLM,
- GenerationConfig,
- PreTrainedTokenizer,
- PreTrainedTokenizerFast,
-)
-from transformers.models.llama.modeling_llama import LlamaForCausalLM
+from diffusers import DiffusionPipeline
+from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
-from colossalai.accelerator import get_accelerator
-from colossalai.cluster import ProcessGroupMesh
-from colossalai.inference.batch_bucket import BatchBucket
-from colossalai.inference.config import InferenceConfig, InputMetaData, ModelShardInferenceConfig
-from colossalai.inference.graph_runner import CUDAGraphRunner
-from colossalai.inference.modeling.policy import model_policy_map
-from colossalai.inference.sampler import search_tokens
-from colossalai.inference.spec import Drafter, GlideInput
-from colossalai.inference.struct import Sequence
-from colossalai.inference.utils import get_model_size, has_index_file
-from colossalai.interface import ModelWrapper
-from colossalai.lazy import LazyInitContext
-from colossalai.logging import get_dist_logger
-from colossalai.pipeline.stage_manager import PipelineStageManager
-from colossalai.shardformer import ShardConfig, ShardFormer
+from colossalai.inference.config import InferenceConfig
+from colossalai.inference.utils import ModelType, get_model_type
from colossalai.shardformer.policies.base_policy import Policy
-from .request_handler import RequestHandler
-
__all__ = ["InferenceEngine"]
-PP_AXIS, TP_AXIS = 0, 1
-
-_supported_models = {
- "LlamaForCausalLM": LlamaForCausalLM,
- "BaichuanForCausalLM": AutoModelForCausalLM,
-}
-
-_BATCH_SIZES_TO_CAPTURE = [1, 2, 4] + [8 * i for i in range(1, 33)]
-
class InferenceEngine:
"""
InferenceEngine which manages the inference process..
Args:
- model_or_path (nn.Module or str): Path or nn.Module of this model.
+ model_or_path (nn.Module or DiffusionPipeline or str): Path or nn.Module or DiffusionPipeline of this model.
tokenizer Optional[(Union[PreTrainedTokenizer, PreTrainedTokenizerFast])]: Path of the tokenizer to use.
inference_config (Optional[InferenceConfig], optional): Store the configuration information related to inference.
verbose (bool): Determine whether or not to log the generation process.
@@ -60,567 +27,68 @@ class InferenceEngine:
def __init__(
self,
- model_or_path: Union[nn.Module, str],
- tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
- inference_config: InferenceConfig,
+ model_or_path: Union[nn.Module, str, DiffusionPipeline],
+ tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast] = None,
+ inference_config: InferenceConfig = None,
verbose: bool = False,
model_policy: Union[Policy, Type[Policy]] = None,
) -> None:
- self.inference_config = inference_config
- self.dtype = inference_config.dtype
- self.high_precision = inference_config.high_precision
+ self.__dict__["_initialized"] = False # use __dict__ directly to avoid calling __setattr__
+ self.model_type = get_model_type(model_or_path=model_or_path)
+ self.engine = None
+ if self.model_type == ModelType.LLM:
+ from .llm_engine import LLMEngine
- self.verbose = verbose
- self.logger = get_dist_logger(__name__)
- self.model_shard_infer_config = inference_config.to_model_shard_inference_config()
+ self.engine = LLMEngine(
+ model_or_path=model_or_path,
+ tokenizer=tokenizer,
+ inference_config=inference_config,
+ verbose=verbose,
+ model_policy=model_policy,
+ )
+ elif self.model_type == ModelType.DIFFUSION_MODEL:
+ from .diffusion_engine import DiffusionEngine
- self.init_model(model_or_path, model_policy, self.model_shard_infer_config)
-
- self.generation_config = inference_config.to_generation_config(self.model_config)
- self.generation_config_dict = self.generation_config.to_dict()
-
- self.tokenizer = tokenizer
- self.tokenizer.pad_token = self.tokenizer.eos_token
-
- self.request_handler = RequestHandler(self.inference_config, self.model_config)
- self.k_cache, self.v_cache = self.request_handler.get_kvcache()
- # DISCUSS maybe move this into batch info?
-
- self.counter = count()
-
- self.use_cuda_graph = self.inference_config.use_cuda_graph
- if self.use_cuda_graph:
- self.graph_runners: Dict[int, CUDAGraphRunner] = {}
- self.graph_memory_pool = None # Set during graph capture.
- if verbose:
- self.logger.info("Colossal AI CUDA Graph Capture on")
-
- self.capture_model(self.k_cache, self.v_cache)
-
- # Model and relatable attrs of speculative decoding will be set by `enable_spec_dec`
- self.use_spec_dec = self.inference_config.use_spec_dec
-
- self.drafter_model = None
- self.drafter = None
- self.use_glide = False
- self.n_spec_tokens = self.inference_config.max_n_spec_tokens
+ self.engine = DiffusionEngine(
+ model_or_path=model_or_path,
+ inference_config=inference_config,
+ verbose=verbose,
+ model_policy=model_policy,
+ )
+ elif self.model_type == ModelType.UNKNOWN:
+ self.logger.error(f"Model Type either Difffusion or LLM!")
+ self._initialized = True
self._verify_args()
- def init_model(
- self,
- model_or_path: Union[nn.Module, str],
- model_policy: Union[Policy, Type[Policy]] = None,
- model_shard_infer_config: ModelShardInferenceConfig = None,
- ):
- """
- Shard model or/and Load weight
-
- Args:
- model_or_path Union[nn.Module, str]: path to the checkpoint or model of transformer format.
- model_policy (Policy): the policy to replace the model.
- model_inference_config: the configuration for modeling initialization when inference.
- model_shard_infer_config (ModelShardInferenceConfig): the configuration for init of module when inference.
- """
- pretrained_path = None
- if isinstance(model_or_path, str):
- import colossalai.interface.pretrained as pretrained_utils
-
- try:
- hf_config = AutoConfig.from_pretrained(model_or_path, trust_remote_code=True, torch_dtype=self.dtype)
- arch = getattr(hf_config, "architectures")[0]
- if arch in _supported_models.keys():
- if arch is "BaichuanForCausalLM":
- self.logger.warning(
- "Attention ! We use lazy init by default, which could be faster for model loading. For baichuan model, the output maybe have a slight difference with transformers"
- )
- ctx = LazyInitContext(default_device="cuda")
- with ctx:
- model = _supported_models[arch].from_pretrained(
- model_or_path, trust_remote_code=True, torch_dtype=self.dtype
- )
- pretrained_path = pretrained_utils.get_pretrained_path(model)
- else:
- # TODO(char-1ee): if the model not supported, use transformers APIs to load and generate
- raise ValueError(f"Model {arch} is not supported.")
-
- except Exception as e:
- self.logger.error(
- f"An exception occurred during loading model: {e}, model should be loaded by transformers\n"
- )
- else:
- model = model_or_path
-
- self.model_config = model.config
-
- torch.cuda.empty_cache()
- init_gpu_memory = torch.cuda.mem_get_info()[0]
-
- self.device = get_accelerator().get_current_device()
- if self.verbose:
- self.logger.info(f"the device is {self.device}")
-
- model = model.to(self.dtype).eval()
-
- if self.verbose:
- self.logger.info(
- f"Before the shard, Rank: [{dist.get_rank()}], model size: {get_model_size(model)} GB, model's device is: {model.device}"
- )
-
- if model_policy is None:
- prefix = "nopadding" if not self.inference_config.pad_input else "padding"
- model_policy_key = f"{prefix}_{getattr(self.model_config, 'model_type', None)}"
- model_policy = model_policy_map.get(model_policy_key)
-
- if not isinstance(model_policy, Policy):
- try:
- model_policy = model_policy()
- except Exception as e:
- raise ValueError(f"Unable to instantiate model policy: {e}")
-
- assert isinstance(model_policy, Policy), f"Invalid type of model policy: {type(model_policy)}"
- pg_mesh = ProcessGroupMesh(self.inference_config.pp_size, self.inference_config.tp_size)
- tp_group = pg_mesh.get_group_along_axis(TP_AXIS)
-
- self.model = self._shardformer(
- model,
- model_policy,
- model_shard_infer_config,
- None,
- tp_group=tp_group,
- )
-
- self.model = ModelWrapper(model).to(self.device)
-
- if self.verbose:
- self.logger.info(
- f"After the shard, Rank: [{dist.get_rank()}], model size: {get_model_size(self.model)} GB, model's device is: {model.device}"
- )
-
- if pretrained_path:
- from colossalai.inference.core.plugin import InferCheckpoint_io
-
- cpt_io = InferCheckpoint_io()
- if_has_index_file, model_index_file = has_index_file(pretrained_path)
- assert if_has_index_file, "the model path is invalid"
- cpt_io.load_model(self.model, model_index_file)
-
- free_gpu_memory, _ = torch.cuda.mem_get_info()
- peak_memory = init_gpu_memory - free_gpu_memory
- if self.verbose:
- self.logger.info(
- f"Rank [{dist.get_rank()}], Model Weight Max Occupy {peak_memory / (1024 ** 3)} GB, Model size: {get_model_size(self.model)} GB"
- )
-
- @torch.inference_mode()
- def capture_model(self, k_cache: List[torch.Tensor], v_cache: List[torch.Tensor]):
- assert self.use_cuda_graph, "please turn on the cuda graph"
-
- if self.verbose:
- self.logger.info("Colossal AI CUDA Graph Capture begin")
-
- t_capture_begin = time.perf_counter()
-
- block_size = self.inference_config.block_size
- head_dim = self.model_config.hidden_size // self.model_config.num_attention_heads
-
- # Prepare dummy inputs. These will be reused for all batch sizes.
- max_batch_size = max(_BATCH_SIZES_TO_CAPTURE)
- max_context_len_to_capture = self.inference_config.max_context_len_to_capture
- max_num_blocks = (max_context_len_to_capture + block_size - 1) // block_size
- input_tokens_ids = torch.zeros(max_batch_size, dtype=torch.long).cuda()
- # self.graph_block_tables = np.zeros((max(_BATCH_SIZES_TO_CAPTURE), max_num_blocks), dtype=np.int32)
- self.graph_block_tables = np.full((max(_BATCH_SIZES_TO_CAPTURE), max_num_blocks), -1, dtype=np.int32)
- self.graph_block_tables[:, 0] = np.arange(max_num_blocks, max_num_blocks + max(_BATCH_SIZES_TO_CAPTURE))
- self.graph_block_tables[0, :] = np.arange(
- 0, max_num_blocks
- ) # NOTE this is a hack to insure cuda grpah could capture the fixed cuda kernel grid in flash decoding, to make the first seqlen as the max_seq_len
- block_tables = torch.from_numpy(self.graph_block_tables).cuda()
- output_tensor = torch.zeros(
- (max_batch_size, self.model_config.num_attention_heads * head_dim), dtype=self.dtype, device=self.device
- )
- fd_inter_tensor = self.request_handler.running_bb.fd_inter_tensor
-
- max_num_seqs = self.inference_config.max_batch_size
- batch_size_capture_list = [bs for bs in _BATCH_SIZES_TO_CAPTURE if bs <= max_num_seqs]
- sequence_lengths = torch.ones(max_batch_size, dtype=torch.int).cuda()
- # NOTE this is a hack to insure cuda grpah could capture the fixed cuda kernel grid in flash decoding, to make the first seqlen as the max_seq_len
- sequence_lengths[0] = torch.tensor(
- self.inference_config.max_context_len_to_capture - 1, dtype=torch.int32
- ).cuda()
-
- # NOTE: Capturing the largest batch size first may help reduce the
- # memory usage of CUDA graph.
- for batch_size in reversed(batch_size_capture_list):
- if self.verbose:
- self.logger.info(f"batch size {batch_size} graph capturing")
-
- input_meta_data = InputMetaData(
- block_tables=block_tables[:batch_size],
- sequence_lengths=sequence_lengths[:batch_size],
- fd_inter_tensor=fd_inter_tensor,
- batch_size=batch_size,
- is_prompts=False,
- use_cuda_graph=True,
- high_precision=False,
- kv_seq_len=sequence_lengths[:batch_size].max().item(),
- head_dim=head_dim,
- dtype=self.dtype,
- )
-
- graph_runner = CUDAGraphRunner(self.model)
- graph_runner.capture(
- input_tokens_ids[:batch_size],
- output_tensor[:batch_size],
- input_meta_data,
- k_caches=k_cache,
- v_caches=v_cache,
- memory_pool=self.graph_memory_pool,
- )
- self.graph_memory_pool = graph_runner.graph.pool()
- self.graph_runners[batch_size] = graph_runner
-
- t_capture_end = time.perf_counter()
-
- if self.verbose:
- self.logger.info(f"CUDA Graph capture time: {t_capture_end - t_capture_begin} s")
-
def _verify_args(self) -> None:
"""Verify the input args"""
- if not isinstance(self.inference_config, InferenceConfig):
- raise TypeError("Invalid type of inference config provided.")
- if not isinstance(self.model, nn.Module):
- raise TypeError(f"the model type must be nn.Module, but got {type(self.model)}")
- if not isinstance(self.tokenizer, (PreTrainedTokenizerFast, PreTrainedTokenizer)):
- raise TypeError(
- f"the tokenizer type must be PreTrainedTokenizer or PreTrainedTokenizerFast, but got {type(self.tokenizer)}"
- )
- if isinstance(self.model, ModelWrapper):
- model = self.model.module
- assert (
- model.__class__.__name__ in _supported_models.keys()
- ), f"Model {self.model.__class__.__name__} is not supported."
-
- def _shardformer(
- self,
- model: nn.Module,
- model_policy: Policy,
- model_shard_infer_config: ModelShardInferenceConfig = None,
- stage_manager: PipelineStageManager = None,
- tp_group: ProcessGroupMesh = None,
- ) -> nn.Module:
- """
- Initialize ShardConfig and replace the model with shardformer.
-
- Args:
- model (nn.Module): Path or nn.Module of this model.
- model_policy (Policy): The policy to shardformer model which is determined by the model type.
- stage_manager (PipelineStageManager, optional): Used to manage pipeline stages. Defaults to None.
- tp_group (ProcessGroupMesh, optional): Used to manage the process TP group mesh. Defaults to None.
-
- Returns:
- nn.Module: The model optimized by Shardformer.
- """
-
- shardconfig = ShardConfig(
- tensor_parallel_process_group=tp_group,
- pipeline_stage_manager=stage_manager,
- enable_tensor_parallelism=(self.inference_config.tp_size > 1),
- enable_fused_normalization=False,
- enable_all_optimization=False,
- enable_flash_attention=False,
- enable_jit_fused=False,
- enable_sequence_parallelism=False,
- extra_kwargs={"model_shard_infer_config": model_shard_infer_config},
- )
- shardformer = ShardFormer(shard_config=shardconfig)
- shard_model, _ = shardformer.optimize(model, model_policy)
- return shard_model
-
- def enable_spec_dec(
- self,
- drafter_model: nn.Module = None,
- n_spec_tokens: int = None,
- use_glide_drafter: bool = False,
- ) -> None:
- """Initialize drafter (if it has not yet), and enable Speculative Decoding for subsequent generations.
-
- Args:
- drafter_model (nn.Module): The drafter model (small model) used to speculate tokens.
- If provided, the previous drafter and drafter model, if exist, will be overwritten.
- n_spec_tokens (Optional[int]): The number of tokens to speculate in each round of speculating-verifying.
- If not provided, `max_n_spec_tokens` in InferenceConfig will be used.
- use_glide_drafter (bool): Whether to use glide model for speculative decoding. Defaults to False.
- If True, the drafter model will be replaced by a glide model.
-
- ```python
- ...
- engine = InferenceEngine(model, tokenizer, inference_config)
-
- engine.enable_spec_dec(drafter_model, n_spec_tokens=5)
- engine.generate(...) # Speculative Decoding
-
- engine.disable_spec_dec()
- engine.generate(...) # Normal generation
-
- engine.enable_spec_dec()
- engine.generate(...) # Speculative-Decoding using previously set drafter model and number of spec tokens
- engine.clear_spec_dec()
- ```
- """
-
- if drafter_model is None and self.drafter is None:
- raise ValueError("Drafter not initialized. Please provide a Drafter Model")
- if n_spec_tokens is not None:
- assert 1 < n_spec_tokens <= self.inference_config.max_n_spec_tokens
- self.n_spec_tokens = n_spec_tokens
- if drafter_model is not None:
- assert isinstance(drafter_model, nn.Module)
- # overwrite the drafter, if exists
- self.clear_spec_dec()
- self.drafter_model = drafter_model
- self.drafter = Drafter(
- self.drafter_model,
- self.tokenizer,
- device=self.device,
- dtype=self.dtype,
- )
-
- # check if the provided drafter model is compatible with GLIDE structure
- # when `use_glide_drafter` is set to True
- if (
- use_glide_drafter
- and hasattr(drafter_model, "model")
- and hasattr(drafter_model.model, "layers")
- and hasattr(drafter_model.model.layers[0], "cross_attn")
- ):
- self.use_glide = use_glide_drafter
- elif use_glide_drafter:
- self.logger.warning(
- f"`use_glide_drafter` is provided as {use_glide_drafter}, "
- f"but the provided drafter model is not compatible with GLIDE structure."
- f"Falling back to use the default drafter model (non-GLIDE)."
- )
- self.request_handler.set_spec_dec_mode(self.n_spec_tokens)
- # using speculative decoding for subsequent generations
- self.use_spec_dec = True
-
- def disable_spec_dec(self) -> None:
- """Disable using speculative decoding for subsequent generations."""
- self.request_handler.unset_spec_dec_mode()
- # set back to the maximum number of tokens to speculate
- self.n_spec_tokens = self.inference_config.max_n_spec_tokens
- self.use_glide = False
- self.use_spec_dec = False
-
- def clear_spec_dec(self) -> None:
- """Clear relatable structures of speculative decoding, if exist."""
- if self.use_spec_dec:
- self.disable_spec_dec()
- if self.drafter_model or self.drafter:
- self.drafter_model = None
- self.drafter = None
- torch.cuda.empty_cache()
- self.use_glide = False
- self.use_spec_dec = False
-
- def steps_spec_dec(self) -> List[Sequence]:
- """
- Run Speculative Decoding steps. This is like retrieving a single batch and launch inference
- with many steps of speculating by a drafter model as well as verifying by a main model.
-
- Returns:
- List[Sequence]: finished sequences generated by one step.
- """
- batch = self.request_handler.schedule() # prefill batch
- assert batch.current_batch_size == 1, "Only support bsz 1 for speculative decoding for now."
-
- input_token_ids, output_tensor, input_meta_data = self.prepare_input(batch)
-
- if input_meta_data.use_cuda_graph:
- model_executable = self.graph_runners[input_meta_data.batch_size]
- else:
- model_executable = self.model
-
- # 1. Prefill small model (Drafter) - fill past kv cache for drafter model
- # NOTE For glide drafter models, we won't actually apply glide during prefill stage
- drafter_out = self.drafter.speculate(input_token_ids, 1, None)
- next_token_ids_spec = drafter_out.next_tokens
- drafter_past_key_values = drafter_out.past_key_values
-
- # 2. Prefill main model (Verifier) - fill past kv cache for main model
- logits = model_executable(input_token_ids, output_tensor, input_meta_data, self.k_cache, self.v_cache)
- next_tokens = search_tokens(self.generation_config, logits, batch_token_ids=batch.batch_token_ids)
- # append new inputs to the batch, temporarily
- batch.append_batch_tokens(next_tokens)
- self.request_handler.allocate_batch_spec_dec(batch, 1)
- already_allocated_kv_len = batch.seq_lengths[0].item()
- input_token_ids = batch.get_1D_inputs_spec_dec(1)
-
- finished_sequences = self.request_handler.update()
-
- while True:
- # HACK Retrieve the running batch
- # Using RequestHandler.schedule here will re-allocate same kv cache for the batch
- batch = self.request_handler.running_bb # running batch
- assert batch.current_batch_size == 1, "Only support bsz 1 for speculative decoding for now."
-
- # 3. Decoding - Drafter model speculates `n` tokens
- glide_input = None
- if self.use_glide:
- glide_input = GlideInput(
- batch.get_block_table_tensor(),
- self.k_cache[-1], # use kv cahces of the last layer
- self.v_cache[-1],
- batch.get_sequence_lengths(),
- n_spec_tokens=self.n_spec_tokens,
- )
-
- drafter_out = self.drafter.speculate(
- input_token_ids,
- self.n_spec_tokens,
- drafter_past_key_values,
- glide_input=glide_input,
- )
- next_token_ids_spec = drafter_out.next_tokens
- drafter_past_key_values = drafter_out.past_key_values
- drafter_spec_length = drafter_out.speculated_length
-
- for next_token_id_spec in next_token_ids_spec:
- self.request_handler.append_next_tokens(next_token_id_spec.unsqueeze(0))
- cur_length = batch.seq_lengths[0].item()
- if already_allocated_kv_len < cur_length:
- self.request_handler.allocate_batch_spec_dec(batch, n=cur_length - already_allocated_kv_len)
- already_allocated_kv_len = cur_length
-
- # 4. Decoding - Main model verifies `n` tokens in parallel
- if drafter_spec_length < batch.num_tokens_to_verify:
- batch.set_use_spec_dec(num_tokens_to_verify=drafter_spec_length)
- input_token_ids, output_tensor, input_meta_data = self.prepare_input(batch)
- logits = model_executable(input_token_ids, output_tensor, input_meta_data, self.k_cache, self.v_cache)
-
- next_tokens = search_tokens(self.generation_config, logits, batch_token_ids=batch.batch_token_ids)
-
- # 5. Compare and process the results
- diff_indexes = torch.nonzero(~(next_tokens[:-1] == next_token_ids_spec))
- n_matches = drafter_spec_length if diff_indexes.size(0) == 0 else diff_indexes[0][0].item()
-
- # revoke appended tokens for each Sequence in the current batch
- batch.revoke_batch_tokens(drafter_spec_length - n_matches) # revoke drafted tokens
-
- # append the last correct token generated by the main model
- self.request_handler.append_next_tokens(next_tokens[n_matches].unsqueeze(0))
-
- # trim past key values of the drafter model
- drafter_past_key_values = Drafter.trim_kv_cache(
- drafter_past_key_values, drafter_spec_length - n_matches - 1
- )
-
- # prepare inputs for the next round of speculation
- n = 1 if n_matches < drafter_spec_length else 2
- input_token_ids = batch.get_1D_inputs_spec_dec(n)
-
- self.request_handler.update_batch_finished(batch, generation_config=self.generation_config)
- finished_sequences = self.request_handler.update()
- if len(finished_sequences) > 0:
- break
-
- # Reset back the number of speculated tokens of the batch,
- # this is used to handle the last round of speculation, in which case the number of speculated tokens
- # by the drafter is less than the number of speculated tokens set to the engine.
- batch.set_use_spec_dec(num_tokens_to_verify=self.n_spec_tokens)
-
- return finished_sequences
+ assert self.engine is not None, "Please init Engine first"
+ assert self._initialized, "Engine must be initialized"
def generate(
self,
request_ids: Union[List[int], int] = None,
prompts: Union[List[str], str] = None,
- prompts_token_ids: Union[List[int], torch.Tensor, np.ndarray] = None,
- return_token_ids: bool = False,
- generation_config: Optional[GenerationConfig] = None,
- ) -> Union[List[str], Tuple[List[str], List[List[int]]]]:
+ *args,
+ **kwargs,
+ ) -> Union[List[Union[str, List[PIL.Image.Image], np.ndarray]], Tuple[List[str], List[List[int]]]]:
"""
Executing the inference step.
Args:
request_ids (List[int], optional): The request ID. Defaults to None.
prompts (Union[List[str], optional): Input prompts. Defaults to None.
- prompts_token_ids (Union[List[int], torch.Tensor, np.ndarray], optional): token ids of input prompts. Defaults to None.
- return_token_ids (bool, optional): Whether to return output token ids. Defaults to False.
- generation_config (Optional[GenerationConfig], optional): Huggingface GenerationConfig used for inference. Defaults to None.
-
- Returns:
- Union[List[str], Tuple[List[str], List[List[int]]]]: Inference result returned by one generation.
"""
- gen_config_dict = generation_config.to_dict() if generation_config is not None else {}
- prompts = [prompts] if isinstance(prompts, str) else prompts
- request_ids = [request_ids] if isinstance(request_ids, int) else request_ids
-
- with torch.inference_mode():
- if prompts is not None or prompts_token_ids is not None:
- self.add_request(
- request_ids=request_ids,
- prompts=prompts,
- prompts_token_ids=prompts_token_ids,
- **gen_config_dict,
- )
-
- output_seqs_list = []
- total_tokens_list = []
-
- # intuition: If user provide a generation config, we should replace the existing one.
- if generation_config is not None:
- self.generation_config = generation_config
- self.generation_config_dict = gen_config_dict
-
- if self.use_spec_dec:
- assert self.drafter is not None, "Drafter Model is not initialized."
- while self.request_handler.check_unfinished_seqs():
- output_seqs_list += self.steps_spec_dec()
- else:
- while self.request_handler.check_unfinished_seqs():
- output_seqs_list += self.step()
-
- output_seqs_list = sorted(output_seqs_list, key=lambda x: int(x.request_id))
-
- for seq in output_seqs_list:
- total_tokens_list.append(seq.input_token_id + seq.output_token_id)
-
- output_str = self.tokenizer.batch_decode(total_tokens_list, skip_special_tokens=True)
-
- if return_token_ids:
- output_tokens_list = [seq.output_token_id for seq in output_seqs_list]
- return output_str, output_tokens_list
- else:
- return output_str
-
- @property
- def has_prompt_template(self) -> bool:
- """ """
- return self.inference_config.prompt_template is not None
-
- def format_prompt(self, prompts: Union[List[str], str]) -> Union[List[str], str]:
- """
- This method will format the input prompt according to the prompt template given to the InferenceConfig.
- """
- assert (
- self.has_prompt_template
- ), "Found the prompt_template is None. Please provide a valid prompt_template in InferenceConfig."
-
- if isinstance(prompts, (list, tuple)):
- return [self.inference_config.prompt_template.format(input_text=prompt) for prompt in prompts]
- elif isinstance(prompts, str):
- return self.inference_config.prompt_template.format(input_text=prompts)
- else:
- raise TypeError(f"Expected the input prompt to be one of list, tuple, or str, but got {type(prompts)}.")
+ assert self.engine is not None, "Please init Engine first"
+ return self.engine.generate(request_ids=request_ids, prompts=prompts, *args, **kwargs)
def add_request(
self,
request_ids: Union[List[int], int] = None,
prompts: Union[List[str], str] = None,
- prompts_token_ids: Union[List[int], torch.Tensor, np.ndarray] = None,
+ *args,
**kwargs,
) -> None:
"""
@@ -630,168 +98,36 @@ class InferenceEngine:
request_ids (List[int], optional): The request ID. Defaults to None.
prompts (Union[List[str], optional): Input prompts. Defaults to None.
prompts_token_ids (List[List[int]], optional): token ids of input prompts. Defaults to None.
+ kwargs: for LLM, it could be max_length, max_new_tokens, etc
+ for diffusion, it could be prompt_2, prompt_3, num_images_per_prompt, do_classifier_free_guidance, negative_prompt, negative_prompt_2, negative_prompt_3, prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds, clip_skip, which aligns with diffusers
"""
+ assert self.engine is not None, "Please init Engine first"
+ self.engine.add_request(request_ids=request_ids, prompts=prompts, *args, **kwargs)
- # apply the prompt template to the input prompts
+ def step(self):
+ assert self.engine is not None, "Please init Engine first"
+ return self.engine.step()
- if self.has_prompt_template and prompts is not None:
- prompts = self.format_prompt(prompts)
-
- block_size = self.inference_config.block_size
-
- if request_ids is not None and not isinstance(request_ids, list):
- request_ids = [request_ids]
-
- if prompts is not None and not isinstance(prompts, list):
- prompts = [prompts]
-
- if prompts_token_ids is None:
- assert prompts, "When the prompts_token_ids is none, the input prompt list must be provided."
- prompts_token_ids = self.tokenizer.batch_encode_plus(prompts, padding=self.inference_config.pad_input)[
- "input_ids"
- ]
-
- # list of torch Tensor
- if isinstance(prompts_token_ids, list):
- if isinstance(prompts_token_ids[0], torch.Tensor):
- prompts_token_ids = [prompt_token_id.tolist() for prompt_token_id in prompts_token_ids]
- elif isinstance(prompts_token_ids, torch.Tensor) or isinstance(prompts_token_ids, np.ndarray):
- prompts_token_ids = prompts_token_ids.tolist()
- else:
- raise TypeError(
- f"The dtype of prompts_token_ids must be one of list, torch.Tensor, np.ndarray, but got {type(prompts_token_ids)}."
- )
-
- assert (
- len(prompts_token_ids[0]) <= self.inference_config.max_input_len
- ), f"The length of input prompts {len(prompts_token_ids[0])} must be less than max_input_len {self.inference_config.max_input_len}."
-
- prompts_num = len(prompts_token_ids)
-
- for i in range(prompts_num):
- if request_ids:
- assert isinstance(
- request_ids[0], int
- ), f"The request_id type must be int, but got {type(request_ids[0])}"
- assert len(request_ids) == prompts_num
- request_id = request_ids[i]
+ def __getattr__(self, name):
+ """
+ The Design logic of getattr, setattr:
+ 1. Since InferenceEngine is a wrapper for DiffusionEngine/LLMEngine, we hope to invoke all the member of DiffusionEngine/LLMEngine like we just call the member of InferenceEngine.
+ 2. When we call the __init__ of InferenceEngine, we don't want to setattr using self.__dict__["xxx"] = xxx, we want to use origin ways like self.xxx = xxx
+ So we set the attribute `_initialized`. And after initialized, if we couldn't get the member from InferenceEngine, we will try to get the member from self.engine(DiffusionEngine/LLMEngine)
+ """
+ if self.__dict__.get("_initialized", False):
+ if name in self.__dict__:
+ return self.__dict__[name]
else:
- request_id = next(self.counter)
- if prompts == None:
- prompt = None
+ return getattr(self.engine, name)
+ else:
+ return self.__dict__[name]
+
+ def __setattr__(self, name, value):
+ if self.__dict__.get("_initialized", False):
+ if name in self.__dict__:
+ self.__dict__[name] = value
else:
- prompt = prompts[i]
-
- max_length = kwargs.get("max_length", None)
- max_new_tokens = kwargs.get("max_new_tokens", None)
- if max_length is None and max_new_tokens is None:
- max_new_tokens = self.generation_config.max_new_tokens or self.inference_config.max_output_len
- elif max_length is not None:
- max_new_tokens = max_length - len(prompts_token_ids[i])
-
- if not self.inference_config.enable_streamingllm:
- assert (
- self.inference_config.max_output_len >= max_new_tokens
- ), f"max_new_tokens={max_new_tokens} must be less than max_output_len={self.inference_config.max_output_len}."
-
- sequence = Sequence(
- request_id,
- prompt,
- prompts_token_ids[i],
- block_size,
- None,
- self.tokenizer.eos_token_id,
- self.tokenizer.pad_token_id,
- max_output_len=max_new_tokens,
- ignore_eos=self.inference_config.ignore_eos,
- )
- self.request_handler.add_sequence(sequence)
-
- def prepare_input(self, batch: BatchBucket) -> Tuple[torch.Tensor, torch.Tensor, InputMetaData]:
- input_ids = batch.get_1D_inputs()
- sequence_lengths = batch.get_sequence_lengths()
-
- if batch.is_prompts:
- n_tokens = sequence_lengths.sum().item()
+ setattr(self.engine, name, value)
else:
- n_tokens = batch.current_batch_size
- if batch.use_spec_dec:
- n_tokens = batch.num_tokens_to_verify + 1
- assert n_tokens == input_ids.size(0)
- n_tokens = n_tokens * batch.current_batch_size
- output_tensor = torch.zeros(
- (n_tokens, batch.num_heads * batch.head_dim), dtype=batch.dtype, device=batch.device
- )
-
- batch_token_ids = None
- if (
- self.generation_config.repetition_penalty != 1.0
- or self.generation_config.no_repeat_ngram_size > 0
- or self.generation_config.forced_eos_token_id is not None
- ):
- batch_token_ids = batch.batch_token_ids
-
- # only when we have the graph for specific decoding batch size can we use the cuda graph for inference
- use_cuda_graph = False
- if self.use_cuda_graph and not batch.is_prompts and batch.current_batch_size in self.graph_runners.keys():
- use_cuda_graph = True
-
- input_meta_data = InputMetaData(
- block_tables=batch.get_block_table_tensor(),
- sequence_lengths=sequence_lengths,
- fd_inter_tensor=batch.fd_inter_tensor,
- batch_size=batch.current_batch_size,
- is_prompts=batch.is_prompts,
- use_cuda_kernel=self.inference_config.use_cuda_kernel,
- use_cuda_graph=use_cuda_graph,
- high_precision=self.high_precision,
- kv_seq_len=sequence_lengths.max().item(),
- head_dim=batch.head_dim,
- dtype=batch.dtype,
- use_spec_dec=batch.use_spec_dec,
- num_tokens_to_verify=batch.num_tokens_to_verify,
- batch_token_ids=batch_token_ids,
- )
-
- return input_ids, output_tensor, input_meta_data
-
- def step(self) -> List[str]:
- """
- In each step, do the follows:
- 1. Run RequestHandler.schedule() and get the batch used for inference.
- 2. Get the input, inputinfo and output placeholder from the batchbucket
- 3. Run model to generate the next token
- 4. Update waiting list and running list in RequestHandler and get finished sequences.
- 5. Decode and return finished sequences.
-
- Returns:
- List[str]: Decoded finished sequences generated by one step.
- """
-
- batch = self.request_handler.schedule()
-
- input_token_ids, output_tensor, input_meta_data = self.prepare_input(batch)
-
- if input_meta_data.use_cuda_graph:
- model_executable = self.graph_runners[input_meta_data.batch_size]
- else:
- model_executable = self.model
-
- # TODO: padding_id is used for generating attn_mask and will be removed if nopad version is supported.
- logits = model_executable(input_token_ids, output_tensor, input_meta_data, self.k_cache, self.v_cache)
- if self.inference_config.pad_input:
- logits = logits[:, -1, :]
-
- if self.inference_config.enable_streamingllm:
- updated_block_ids = batch.streamingllm_update_batch(
- self.inference_config.start_token_size, self.inference_config.generated_token_size
- )
- self.request_handler.streamingllm_free_block_tables(updated_block_ids)
-
- next_tokens = search_tokens(
- self.generation_config, logits, input_meta_data.is_prompts, batch_token_ids=input_meta_data.batch_token_ids
- )
- self.request_handler.append_next_tokens(next_tokens)
- finished_sequences = self.request_handler.update()
-
- return finished_sequences
+ self.__dict__[name] = value
diff --git a/colossalai/inference/core/llm_engine.py b/colossalai/inference/core/llm_engine.py
new file mode 100644
index 000000000..1dbc3ace8
--- /dev/null
+++ b/colossalai/inference/core/llm_engine.py
@@ -0,0 +1,758 @@
+import time
+from itertools import count
+from typing import Dict, List, Optional, Tuple, Type, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch import distributed as dist
+from transformers import (
+ AutoConfig,
+ AutoModelForCausalLM,
+ GenerationConfig,
+ PreTrainedTokenizer,
+ PreTrainedTokenizerFast,
+)
+from transformers.models.llama.modeling_llama import LlamaForCausalLM
+
+from colossalai.accelerator import get_accelerator
+from colossalai.cluster import ProcessGroupMesh
+from colossalai.inference.batch_bucket import BatchBucket
+from colossalai.inference.config import InferenceConfig, InputMetaData, ModelShardInferenceConfig
+from colossalai.inference.graph_runner import CUDAGraphRunner
+from colossalai.inference.modeling.policy import model_policy_map
+from colossalai.inference.sampler import search_tokens
+from colossalai.inference.spec import Drafter, GlideInput
+from colossalai.inference.struct import Sequence
+from colossalai.inference.utils import get_model_size, has_index_file
+from colossalai.interface import ModelWrapper
+from colossalai.lazy import LazyInitContext
+from colossalai.logging import get_dist_logger
+from colossalai.shardformer.policies.base_policy import Policy
+
+from .base_engine import BaseEngine
+from .request_handler import RequestHandler
+
+PP_AXIS, TP_AXIS = 0, 1
+
+_supported_models = {
+ "LlamaForCausalLM": LlamaForCausalLM,
+ "BaichuanForCausalLM": AutoModelForCausalLM,
+}
+
+_BATCH_SIZES_TO_CAPTURE = [1, 2, 4] + [8 * i for i in range(1, 33)]
+
+
+class LLMEngine(BaseEngine):
+ """
+ InferenceEngine which manages the inference process..
+
+ Args:
+ model_or_path (nn.Module or str): Path or nn.Module of this model.
+ tokenizer Optional[(Union[PreTrainedTokenizer, PreTrainedTokenizerFast])]: Path of the tokenizer to use.
+ inference_config (Optional[InferenceConfig], optional): Store the configuration information related to inference.
+ verbose (bool): Determine whether or not to log the generation process.
+ model_policy ("Policy"): the policy to shardformer model. It will be determined by the model type if not provided.
+ """
+
+ def __init__(
+ self,
+ model_or_path: Union[nn.Module, str],
+ tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast] = None,
+ inference_config: InferenceConfig = None,
+ verbose: bool = False,
+ model_policy: Union[Policy, type[Policy]] = None,
+ ) -> None:
+ self.inference_config = inference_config
+ self.dtype = inference_config.dtype
+ self.high_precision = inference_config.high_precision
+
+ self.verbose = verbose
+ self.logger = get_dist_logger(__name__)
+ self.model_shard_infer_config = inference_config.to_model_shard_inference_config()
+
+ self.init_model(model_or_path, model_policy, self.model_shard_infer_config)
+
+ self.generation_config = inference_config.to_generation_config(self.model_config)
+ self.generation_config_dict = self.generation_config.to_dict()
+
+ self.tokenizer = tokenizer
+ self.tokenizer.pad_token = self.tokenizer.eos_token
+
+ self.request_handler = RequestHandler(self.inference_config, self.model_config)
+ self.k_cache, self.v_cache = self.request_handler.get_kvcache()
+ # DISCUSS maybe move this into batch info?
+
+ self.counter = count()
+
+ self.use_cuda_graph = self.inference_config.use_cuda_graph
+ if self.use_cuda_graph:
+ self.graph_runners: Dict[int, CUDAGraphRunner] = {}
+ self.graph_memory_pool = None # Set during graph capture.
+ if verbose:
+ self.logger.info("Colossal AI CUDA Graph Capture on")
+
+ self.capture_model(self.k_cache, self.v_cache)
+
+ # Model and relatable attrs of speculative decoding will be set by `enable_spec_dec`
+ self.use_spec_dec = self.inference_config.use_spec_dec
+
+ self.drafter_model = None
+ self.drafter = None
+ self.use_glide = False
+ self.n_spec_tokens = self.inference_config.max_n_spec_tokens
+
+ self._verify_args()
+
+ def init_model(
+ self,
+ model_or_path: Union[nn.Module, str],
+ model_policy: Union[Policy, Type[Policy]] = None,
+ model_shard_infer_config: ModelShardInferenceConfig = None,
+ ):
+ """
+ Shard model or/and Load weight
+
+ Args:
+ model_or_path Union[nn.Module, str]: path to the checkpoint or model of transformer format.
+ model_policy (Policy): the policy to replace the model.
+ model_inference_config: the configuration for modeling initialization when inference.
+ model_shard_infer_config (ModelShardInferenceConfig): the configuration for init of module when inference.
+ """
+ pretrained_path = None
+ if isinstance(model_or_path, str):
+ import colossalai.interface.pretrained as pretrained_utils
+
+ try:
+ hf_config = AutoConfig.from_pretrained(model_or_path, trust_remote_code=True, torch_dtype=self.dtype)
+ arch = getattr(hf_config, "architectures")[0]
+ if arch in _supported_models.keys():
+ if arch == "BaichuanForCausalLM":
+ self.logger.warning(
+ "Attention ! We use lazy init by default, which could be faster for model loading. For baichuan model, the output maybe have a slight difference with transformers"
+ )
+ ctx = LazyInitContext(default_device="cuda")
+ with ctx:
+ model = _supported_models[arch].from_pretrained(
+ model_or_path, trust_remote_code=True, torch_dtype=self.dtype
+ )
+ pretrained_path = pretrained_utils.get_pretrained_path(model)
+ else:
+ # TODO(char-1ee): if the model not supported, use transformers APIs to load and generate
+ raise ValueError(f"Model {arch} is not supported.")
+
+ except Exception as e:
+ self.logger.error(
+ f"An exception occurred during loading model: {e}, model should be loaded by transformers\n"
+ )
+ else:
+ model = model_or_path
+
+ self.model_config = model.config
+
+ torch.cuda.empty_cache()
+ init_gpu_memory = torch.cuda.mem_get_info()[0]
+
+ self.device = get_accelerator().get_current_device()
+ if self.verbose:
+ self.logger.info(f"the device is {self.device}")
+
+ model = model.to(self.dtype).eval()
+
+ if self.verbose:
+ self.logger.info(
+ f"Before the shard, Rank: [{dist.get_rank()}], model size: {get_model_size(model)} GB, model's device is: {model.device}"
+ )
+
+ if model_policy is None:
+ prefix = "nopadding" if not self.inference_config.pad_input else "padding"
+ model_policy_key = f"{prefix}_{getattr(self.model_config, 'model_type', None)}"
+ model_policy = model_policy_map.get(model_policy_key)
+
+ if not isinstance(model_policy, Policy):
+ try:
+ model_policy = model_policy()
+ except Exception as e:
+ raise ValueError(f"Unable to instantiate model policy: {e}")
+
+ assert isinstance(model_policy, Policy), f"Invalid type of model policy: {type(model_policy)}"
+ pg_mesh = ProcessGroupMesh(self.inference_config.pp_size, self.inference_config.tp_size)
+ tp_group = pg_mesh.get_group_along_axis(TP_AXIS)
+
+ self.model = self._shardformer(
+ model,
+ model_policy,
+ model_shard_infer_config,
+ None,
+ tp_group=tp_group,
+ )
+
+ self.model = ModelWrapper(model).to(self.device)
+
+ if self.verbose:
+ self.logger.info(
+ f"After the shard, Rank: [{dist.get_rank()}], model size: {get_model_size(self.model)} GB, model's device is: {model.device}"
+ )
+
+ if pretrained_path:
+ from colossalai.inference.core.plugin import InferCheckpoint_io
+
+ cpt_io = InferCheckpoint_io()
+ if_has_index_file, model_index_file = has_index_file(pretrained_path)
+ assert if_has_index_file, "the model path is invalid"
+ cpt_io.load_model(self.model, model_index_file)
+
+ free_gpu_memory, _ = torch.cuda.mem_get_info()
+ peak_memory = init_gpu_memory - free_gpu_memory
+ if self.verbose:
+ self.logger.info(
+ f"Rank [{dist.get_rank()}], Model Weight Max Occupy {peak_memory / (1024 ** 3)} GB, Model size: {get_model_size(self.model)} GB"
+ )
+
+ @torch.inference_mode()
+ def capture_model(self, k_cache: List[torch.Tensor], v_cache: List[torch.Tensor]):
+ assert self.use_cuda_graph, "please turn on the cuda graph"
+
+ if self.verbose:
+ self.logger.info("Colossal AI CUDA Graph Capture begin")
+
+ t_capture_begin = time.perf_counter()
+
+ block_size = self.inference_config.block_size
+ head_dim = self.model_config.hidden_size // self.model_config.num_attention_heads
+
+ # Prepare dummy inputs. These will be reused for all batch sizes.
+ max_batch_size = max(_BATCH_SIZES_TO_CAPTURE)
+ max_context_len_to_capture = self.inference_config.max_context_len_to_capture
+ max_num_blocks = (max_context_len_to_capture + block_size - 1) // block_size
+ input_tokens_ids = torch.zeros(max_batch_size, dtype=torch.long).cuda()
+ # self.graph_block_tables = np.zeros((max(_BATCH_SIZES_TO_CAPTURE), max_num_blocks), dtype=np.int32)
+ self.graph_block_tables = np.full((max(_BATCH_SIZES_TO_CAPTURE), max_num_blocks), -1, dtype=np.int32)
+ self.graph_block_tables[:, 0] = np.arange(max_num_blocks, max_num_blocks + max(_BATCH_SIZES_TO_CAPTURE))
+ self.graph_block_tables[0, :] = np.arange(
+ 0, max_num_blocks
+ ) # NOTE this is a hack to insure cuda grpah could capture the fixed cuda kernel grid in flash decoding, to make the first seqlen as the max_seq_len
+ block_tables = torch.from_numpy(self.graph_block_tables).cuda()
+ output_tensor = torch.zeros(
+ (max_batch_size, self.model_config.num_attention_heads * head_dim), dtype=self.dtype, device=self.device
+ )
+ fd_inter_tensor = self.request_handler.running_bb.fd_inter_tensor
+
+ max_num_seqs = self.inference_config.max_batch_size
+ batch_size_capture_list = [bs for bs in _BATCH_SIZES_TO_CAPTURE if bs <= max_num_seqs]
+ sequence_lengths = torch.ones(max_batch_size, dtype=torch.int).cuda()
+ # NOTE this is a hack to insure cuda grpah could capture the fixed cuda kernel grid in flash decoding, to make the first seqlen as the max_seq_len
+ sequence_lengths[0] = torch.tensor(
+ self.inference_config.max_context_len_to_capture - 1, dtype=torch.int32
+ ).cuda()
+
+ # NOTE: Capturing the largest batch size first may help reduce the
+ # memory usage of CUDA graph.
+ for batch_size in reversed(batch_size_capture_list):
+ if self.verbose:
+ self.logger.info(f"batch size {batch_size} graph capturing")
+
+ input_meta_data = InputMetaData(
+ block_tables=block_tables[:batch_size],
+ sequence_lengths=sequence_lengths[:batch_size],
+ fd_inter_tensor=fd_inter_tensor,
+ batch_size=batch_size,
+ is_prompts=False,
+ use_cuda_graph=True,
+ high_precision=False,
+ kv_seq_len=sequence_lengths[:batch_size].max().item(),
+ head_dim=head_dim,
+ dtype=self.dtype,
+ )
+
+ graph_runner = CUDAGraphRunner(self.model)
+ graph_runner.capture(
+ input_tokens_ids[:batch_size],
+ output_tensor[:batch_size],
+ input_meta_data,
+ k_caches=k_cache,
+ v_caches=v_cache,
+ memory_pool=self.graph_memory_pool,
+ )
+ self.graph_memory_pool = graph_runner.graph.pool()
+ self.graph_runners[batch_size] = graph_runner
+
+ t_capture_end = time.perf_counter()
+
+ if self.verbose:
+ self.logger.info(f"CUDA Graph capture time: {t_capture_end - t_capture_begin} s")
+
+ def _verify_args(self) -> None:
+ """Verify the input args"""
+ if not isinstance(self.inference_config, InferenceConfig):
+ raise TypeError("Invalid type of inference config provided.")
+ if not isinstance(self.model, nn.Module):
+ raise TypeError(f"the model type must be nn.Module, but got {type(self.model)}")
+ if not isinstance(self.tokenizer, (PreTrainedTokenizerFast, PreTrainedTokenizer)):
+ raise TypeError(
+ f"the tokenizer type must be PreTrainedTokenizer or PreTrainedTokenizerFast, but got {type(self.tokenizer)}"
+ )
+ if isinstance(self.model, ModelWrapper):
+ model = self.model.module
+ assert (
+ model.__class__.__name__ in _supported_models.keys()
+ ), f"Model {self.model.__class__.__name__} is not supported."
+
+ def enable_spec_dec(
+ self,
+ drafter_model: nn.Module = None,
+ n_spec_tokens: int = None,
+ use_glide_drafter: bool = False,
+ ) -> None:
+ """Initialize drafter (if it has not yet), and enable Speculative Decoding for subsequent generations.
+
+ Args:
+ drafter_model (nn.Module): The drafter model (small model) used to speculate tokens.
+ If provided, the previous drafter and drafter model, if exist, will be overwritten.
+ n_spec_tokens (Optional[int]): The number of tokens to speculate in each round of speculating-verifying.
+ If not provided, `max_n_spec_tokens` in InferenceConfig will be used.
+ use_glide_drafter (bool): Whether to use glide model for speculative decoding. Defaults to False.
+ If True, the drafter model will be replaced by a glide model.
+
+ ```python
+ ...
+ engine = InferenceEngine(model, tokenizer, inference_config)
+
+ engine.enable_spec_dec(drafter_model, n_spec_tokens=5)
+ engine.generate(...) # Speculative Decoding
+
+ engine.disable_spec_dec()
+ engine.generate(...) # Normal generation
+
+ engine.enable_spec_dec()
+ engine.generate(...) # Speculative-Decoding using previously set drafter model and number of spec tokens
+ engine.clear_spec_dec()
+ ```
+ """
+
+ if drafter_model is None and self.drafter is None:
+ raise ValueError("Drafter not initialized. Please provide a Drafter Model")
+ if n_spec_tokens is not None:
+ assert 1 < n_spec_tokens <= self.inference_config.max_n_spec_tokens
+ self.n_spec_tokens = n_spec_tokens
+ if drafter_model is not None:
+ assert isinstance(drafter_model, nn.Module)
+ # overwrite the drafter, if exists
+ self.clear_spec_dec()
+ self.drafter_model = drafter_model
+ self.drafter = Drafter(
+ self.drafter_model,
+ self.tokenizer,
+ device=self.device,
+ dtype=self.dtype,
+ )
+
+ # check if the provided drafter model is compatible with GLIDE structure
+ # when `use_glide_drafter` is set to True
+ if (
+ use_glide_drafter
+ and hasattr(drafter_model, "model")
+ and hasattr(drafter_model.model, "layers")
+ and hasattr(drafter_model.model.layers[0], "cross_attn")
+ ):
+ self.use_glide = use_glide_drafter
+ elif use_glide_drafter:
+ self.logger.warning(
+ f"`use_glide_drafter` is provided as {use_glide_drafter}, "
+ f"but the provided drafter model is not compatible with GLIDE structure."
+ f"Falling back to use the default drafter model (non-GLIDE)."
+ )
+ self.request_handler.set_spec_dec_mode(self.n_spec_tokens)
+ # using speculative decoding for subsequent generations
+ self.use_spec_dec = True
+
+ def disable_spec_dec(self) -> None:
+ """Disable using speculative decoding for subsequent generations."""
+ self.request_handler.unset_spec_dec_mode()
+ # set back to the maximum number of tokens to speculate
+ self.n_spec_tokens = self.inference_config.max_n_spec_tokens
+ self.use_glide = False
+ self.use_spec_dec = False
+
+ def clear_spec_dec(self) -> None:
+ """Clear relatable structures of speculative decoding, if exist."""
+ if self.use_spec_dec:
+ self.disable_spec_dec()
+ if self.drafter_model or self.drafter:
+ self.drafter_model = None
+ self.drafter = None
+ torch.cuda.empty_cache()
+ self.use_glide = False
+ self.use_spec_dec = False
+
+ def steps_spec_dec(self) -> List[Sequence]:
+ """
+ Run Speculative Decoding steps. This is like retrieving a single batch and launch inference
+ with many steps of speculating by a drafter model as well as verifying by a main model.
+
+ Returns:
+ List[Sequence]: finished sequences generated by one step.
+ """
+ batch = self.request_handler.schedule() # prefill batch
+ assert batch.current_batch_size == 1, "Only support bsz 1 for speculative decoding for now."
+
+ input_token_ids, output_tensor, input_meta_data = self.prepare_input(batch)
+
+ if input_meta_data.use_cuda_graph:
+ model_executable = self.graph_runners[input_meta_data.batch_size]
+ else:
+ model_executable = self.model
+
+ # 1. Prefill small model (Drafter) - fill past kv cache for drafter model
+ # NOTE For glide drafter models, we won't actually apply glide during prefill stage
+ drafter_out = self.drafter.speculate(input_token_ids, 1, None)
+ next_token_ids_spec = drafter_out.next_tokens
+ drafter_past_key_values = drafter_out.past_key_values
+
+ # 2. Prefill main model (Verifier) - fill past kv cache for main model
+ logits = model_executable(input_token_ids, output_tensor, input_meta_data, self.k_cache, self.v_cache)
+ next_tokens = search_tokens(self.generation_config, logits, batch_token_ids=batch.batch_token_ids)
+ # append new inputs to the batch, temporarily
+ batch.append_batch_tokens(next_tokens)
+ self.request_handler.allocate_batch_spec_dec(batch, 1)
+ already_allocated_kv_len = batch.seq_lengths[0].item()
+ input_token_ids = batch.get_1D_inputs_spec_dec(1)
+
+ finished_sequences = self.request_handler.update()
+
+ while True:
+ # HACK Retrieve the running batch
+ # Using RequestHandler.schedule here will re-allocate same kv cache for the batch
+ batch = self.request_handler.running_bb # running batch
+ assert batch.current_batch_size == 1, "Only support bsz 1 for speculative decoding for now."
+
+ # 3. Decoding - Drafter model speculates `n` tokens
+ glide_input = None
+ if self.use_glide:
+ glide_input = GlideInput(
+ batch.get_block_table_tensor(),
+ self.k_cache[-1], # use kv cahces of the last layer
+ self.v_cache[-1],
+ batch.get_sequence_lengths(),
+ n_spec_tokens=self.n_spec_tokens,
+ )
+
+ drafter_out = self.drafter.speculate(
+ input_token_ids,
+ self.n_spec_tokens,
+ drafter_past_key_values,
+ glide_input=glide_input,
+ )
+ next_token_ids_spec = drafter_out.next_tokens
+ drafter_past_key_values = drafter_out.past_key_values
+ drafter_spec_length = drafter_out.speculated_length
+
+ for next_token_id_spec in next_token_ids_spec:
+ self.request_handler.append_next_tokens(next_token_id_spec.unsqueeze(0))
+ cur_length = batch.seq_lengths[0].item()
+ if already_allocated_kv_len < cur_length:
+ self.request_handler.allocate_batch_spec_dec(batch, n=cur_length - already_allocated_kv_len)
+ already_allocated_kv_len = cur_length
+
+ # 4. Decoding - Main model verifies `n` tokens in parallel
+ if drafter_spec_length < batch.num_tokens_to_verify:
+ batch.set_use_spec_dec(num_tokens_to_verify=drafter_spec_length)
+ input_token_ids, output_tensor, input_meta_data = self.prepare_input(batch)
+ logits = model_executable(input_token_ids, output_tensor, input_meta_data, self.k_cache, self.v_cache)
+
+ next_tokens = search_tokens(self.generation_config, logits, batch_token_ids=batch.batch_token_ids)
+
+ # 5. Compare and process the results
+ diff_indexes = torch.nonzero(~(next_tokens[:-1] == next_token_ids_spec))
+ n_matches = drafter_spec_length if diff_indexes.size(0) == 0 else diff_indexes[0][0].item()
+
+ # revoke appended tokens for each Sequence in the current batch
+ batch.revoke_batch_tokens(drafter_spec_length - n_matches) # revoke drafted tokens
+
+ # append the last correct token generated by the main model
+ self.request_handler.append_next_tokens(next_tokens[n_matches].unsqueeze(0))
+
+ # trim past key values of the drafter model
+ drafter_past_key_values = Drafter.trim_kv_cache(
+ drafter_past_key_values, drafter_spec_length - n_matches - 1
+ )
+
+ # prepare inputs for the next round of speculation
+ n = 1 if n_matches < drafter_spec_length else 2
+ input_token_ids = batch.get_1D_inputs_spec_dec(n)
+
+ self.request_handler.update_batch_finished(batch, generation_config=self.generation_config)
+ finished_sequences = self.request_handler.update()
+ if len(finished_sequences) > 0:
+ break
+
+ # Reset back the number of speculated tokens of the batch,
+ # this is used to handle the last round of speculation, in which case the number of speculated tokens
+ # by the drafter is less than the number of speculated tokens set to the engine.
+ batch.set_use_spec_dec(num_tokens_to_verify=self.n_spec_tokens)
+
+ return finished_sequences
+
+ def generate(
+ self,
+ request_ids: Union[List[int], int] = None,
+ prompts: Union[List[str], str] = None,
+ prompts_token_ids: Union[List[int], torch.Tensor, np.ndarray] = None,
+ return_token_ids: bool = False,
+ generation_config: Optional[GenerationConfig] = None,
+ ) -> Union[List[str], Tuple[List[str], List[List[int]]]]:
+ """
+ Executing the inference step.
+
+ Args:
+ request_ids (List[int], optional): The request ID. Defaults to None.
+ prompts (Union[List[str], optional): Input prompts. Defaults to None.
+ prompts_token_ids (Union[List[int], torch.Tensor, np.ndarray], optional): token ids of input prompts. Defaults to None.
+ return_token_ids (bool, optional): Whether to return output token ids. Defaults to False.
+ generation_config (Optional[GenerationConfig], optional): Huggingface GenerationConfig used for inference. Defaults to None.
+
+ Returns:
+ Union[List[str], Tuple[List[str], List[List[int]]]]: Inference result returned by one generation.
+ """
+
+ gen_config_dict = generation_config.to_dict() if generation_config is not None else {}
+ prompts = [prompts] if isinstance(prompts, str) else prompts
+ request_ids = [request_ids] if isinstance(request_ids, int) else request_ids
+
+ with torch.inference_mode():
+ if prompts is not None or prompts_token_ids is not None:
+ self.add_request(
+ request_ids=request_ids,
+ prompts=prompts,
+ prompts_token_ids=prompts_token_ids,
+ **gen_config_dict,
+ )
+
+ output_seqs_list = []
+ total_tokens_list = []
+
+ # intuition: If user provide a generation config, we should replace the existing one.
+ if generation_config is not None:
+ self.generation_config = generation_config
+ self.generation_config_dict = gen_config_dict
+
+ if self.use_spec_dec:
+ assert self.drafter is not None, "Drafter Model is not initialized."
+ while self.request_handler.check_unfinished_reqs():
+ output_seqs_list += self.steps_spec_dec()
+ else:
+ while self.request_handler.check_unfinished_reqs():
+ output_seqs_list += self.step()
+
+ output_seqs_list = sorted(output_seqs_list, key=lambda x: int(x.request_id))
+
+ for seq in output_seqs_list:
+ total_tokens_list.append(seq.input_token_id + seq.output_token_id)
+
+ output_str = self.tokenizer.batch_decode(total_tokens_list, skip_special_tokens=True)
+
+ if return_token_ids:
+ output_tokens_list = [seq.output_token_id for seq in output_seqs_list]
+ return output_str, output_tokens_list
+ else:
+ return output_str
+
+ @property
+ def has_prompt_template(self) -> bool:
+ """ """
+ return self.inference_config.prompt_template is not None
+
+ def format_prompt(self, prompts: Union[List[str], str]) -> Union[List[str], str]:
+ """
+ This method will format the input prompt according to the prompt template given to the InferenceConfig.
+ """
+ assert (
+ self.has_prompt_template
+ ), "Found the prompt_template is None. Please provide a valid prompt_template in InferenceConfig."
+
+ if isinstance(prompts, (list, tuple)):
+ return [self.inference_config.prompt_template.format(input_text=prompt) for prompt in prompts]
+ elif isinstance(prompts, str):
+ return self.inference_config.prompt_template.format(input_text=prompts)
+ else:
+ raise TypeError(f"Expected the input prompt to be one of list, tuple, or str, but got {type(prompts)}.")
+
+ def add_request(
+ self,
+ request_ids: Union[List[int], int] = None,
+ prompts: Union[List[str], str] = None,
+ prompts_token_ids: Union[List[int], torch.Tensor, np.ndarray] = None,
+ **kwargs,
+ ) -> None:
+ """
+ Add requests.
+
+ Args:
+ request_ids (List[int], optional): The request ID. Defaults to None.
+ prompts (Union[List[str], optional): Input prompts. Defaults to None.
+ prompts_token_ids (List[List[int]], optional): token ids of input prompts. Defaults to None.
+ """
+
+ # apply the prompt template to the input prompts
+
+ if self.has_prompt_template and prompts is not None:
+ prompts = self.format_prompt(prompts)
+
+ block_size = self.inference_config.block_size
+
+ if request_ids is not None and not isinstance(request_ids, list):
+ request_ids = [request_ids]
+
+ if prompts is not None and not isinstance(prompts, list):
+ prompts = [prompts]
+
+ if prompts_token_ids is None:
+ assert prompts, "When the prompts_token_ids is none, the input prompt list must be provided."
+ prompts_token_ids = self.tokenizer.batch_encode_plus(prompts, padding=self.inference_config.pad_input)[
+ "input_ids"
+ ]
+
+ # list of torch Tensor
+ if isinstance(prompts_token_ids, list):
+ if isinstance(prompts_token_ids[0], torch.Tensor):
+ prompts_token_ids = [prompt_token_id.tolist() for prompt_token_id in prompts_token_ids]
+ elif isinstance(prompts_token_ids, torch.Tensor) or isinstance(prompts_token_ids, np.ndarray):
+ prompts_token_ids = prompts_token_ids.tolist()
+ else:
+ raise TypeError(
+ f"The dtype of prompts_token_ids must be one of list, torch.Tensor, np.ndarray, but got {type(prompts_token_ids)}."
+ )
+
+ assert (
+ len(prompts_token_ids[0]) <= self.inference_config.max_input_len
+ ), f"The length of input prompts {len(prompts_token_ids[0])} must be less than max_input_len {self.inference_config.max_input_len}."
+
+ prompts_num = len(prompts_token_ids)
+
+ for i in range(prompts_num):
+ if request_ids:
+ assert isinstance(
+ request_ids[0], int
+ ), f"The request_id type must be int, but got {type(request_ids[0])}"
+ assert len(request_ids) == prompts_num
+ request_id = request_ids[i]
+ else:
+ request_id = next(self.counter)
+ if prompts == None:
+ prompt = None
+ else:
+ prompt = prompts[i]
+
+ max_length = kwargs.get("max_length", None)
+ max_new_tokens = kwargs.get("max_new_tokens", None)
+ if max_length is None and max_new_tokens is None:
+ max_new_tokens = self.generation_config.max_new_tokens or self.inference_config.max_output_len
+ elif max_length is not None:
+ max_new_tokens = max_length - len(prompts_token_ids[i])
+
+ if not self.inference_config.enable_streamingllm:
+ assert (
+ self.inference_config.max_output_len >= max_new_tokens
+ ), f"max_new_tokens={max_new_tokens} must be less than max_output_len={self.inference_config.max_output_len}."
+
+ sequence = Sequence(
+ request_id,
+ prompt,
+ prompts_token_ids[i],
+ block_size,
+ None,
+ self.tokenizer.eos_token_id,
+ self.tokenizer.pad_token_id,
+ max_output_len=max_new_tokens,
+ ignore_eos=self.inference_config.ignore_eos,
+ )
+ self.request_handler.add_sequence(sequence)
+
+ def prepare_input(self, batch: BatchBucket) -> Tuple[torch.Tensor, torch.Tensor, InputMetaData]:
+ input_ids = batch.get_1D_inputs()
+ sequence_lengths = batch.get_sequence_lengths()
+
+ if batch.is_prompts:
+ n_tokens = sequence_lengths.sum().item()
+ else:
+ n_tokens = batch.current_batch_size
+ if batch.use_spec_dec:
+ n_tokens = batch.num_tokens_to_verify + 1
+ assert n_tokens == input_ids.size(0)
+ n_tokens = n_tokens * batch.current_batch_size
+ output_tensor = torch.zeros(
+ (n_tokens, batch.num_heads * batch.head_dim), dtype=batch.dtype, device=batch.device
+ )
+
+ batch_token_ids = None
+ if (
+ self.generation_config.repetition_penalty != 1.0
+ or self.generation_config.no_repeat_ngram_size > 0
+ or self.generation_config.forced_eos_token_id is not None
+ ):
+ batch_token_ids = batch.batch_token_ids
+
+ # only when we have the graph for specific decoding batch size can we use the cuda graph for inference
+ use_cuda_graph = False
+ if self.use_cuda_graph and not batch.is_prompts and batch.current_batch_size in self.graph_runners.keys():
+ use_cuda_graph = True
+
+ input_meta_data = InputMetaData(
+ block_tables=batch.get_block_table_tensor(),
+ sequence_lengths=sequence_lengths,
+ fd_inter_tensor=batch.fd_inter_tensor,
+ batch_size=batch.current_batch_size,
+ is_prompts=batch.is_prompts,
+ use_cuda_kernel=self.inference_config.use_cuda_kernel,
+ use_cuda_graph=use_cuda_graph,
+ high_precision=self.high_precision,
+ kv_seq_len=sequence_lengths.max().item(),
+ head_dim=batch.head_dim,
+ dtype=batch.dtype,
+ use_spec_dec=batch.use_spec_dec,
+ num_tokens_to_verify=batch.num_tokens_to_verify,
+ batch_token_ids=batch_token_ids,
+ )
+
+ return input_ids, output_tensor, input_meta_data
+
+ def step(self) -> List[str]:
+ """
+ In each step, do the follows:
+ 1. Run RequestHandler.schedule() and get the batch used for inference.
+ 2. Get the input, inputinfo and output placeholder from the batchbucket
+ 3. Run model to generate the next token
+ 4. Update waiting list and running list in RequestHandler and get finished sequences.
+ 5. Decode and return finished sequences.
+
+ Returns:
+ List[str]: Decoded finished sequences generated by one step.
+ """
+
+ batch = self.request_handler.schedule()
+
+ input_token_ids, output_tensor, input_meta_data = self.prepare_input(batch)
+
+ if input_meta_data.use_cuda_graph:
+ model_executable = self.graph_runners[input_meta_data.batch_size]
+ else:
+ model_executable = self.model
+
+ # TODO: padding_id is used for generating attn_mask and will be removed if nopad version is supported.
+ logits = model_executable(input_token_ids, output_tensor, input_meta_data, self.k_cache, self.v_cache)
+ if self.inference_config.pad_input:
+ logits = logits[:, -1, :]
+
+ if self.inference_config.enable_streamingllm:
+ updated_block_ids = batch.streamingllm_update_batch(
+ self.inference_config.start_token_size, self.inference_config.generated_token_size
+ )
+ self.request_handler.streamingllm_free_block_tables(updated_block_ids)
+
+ next_tokens = search_tokens(
+ self.generation_config, logits, input_meta_data.is_prompts, batch_token_ids=input_meta_data.batch_token_ids
+ )
+ self.request_handler.append_next_tokens(next_tokens)
+ finished_sequences = self.request_handler.update()
+
+ return finished_sequences
diff --git a/colossalai/inference/core/request_handler.py b/colossalai/inference/core/request_handler.py
index 512eaea71..393347c31 100644
--- a/colossalai/inference/core/request_handler.py
+++ b/colossalai/inference/core/request_handler.py
@@ -8,7 +8,7 @@ from colossalai.inference.batch_bucket import BatchBucket
from colossalai.inference.config import InferenceConfig
from colossalai.inference.flash_decoding_utils import FDIntermTensors
from colossalai.inference.kv_cache import KVCacheManager, RPCKVCacheManager
-from colossalai.inference.struct import RequestStatus, Sequence
+from colossalai.inference.struct import DiffusionSequence, RequestStatus, Sequence
from colossalai.logging import get_dist_logger
logger = get_dist_logger(__name__)
@@ -98,7 +98,46 @@ class RunningList:
self._decoding[seq_id] = self._prefill.pop(seq_id)
-class RequestHandler:
+class NaiveRequestHandler:
+ def __init__(self) -> None:
+ self.running_list: List[DiffusionSequence] = []
+ self.waiting_list: List[str] = []
+
+ def _has_waiting(self) -> bool:
+ return any(lst for lst in self.waiting_list)
+
+ def _has_running(self) -> bool:
+ return any(lst for lst in self.running_list)
+
+ def check_unfinished_reqs(self):
+ return self._has_waiting() or self._has_running()
+
+ def add_sequence(self, seq: DiffusionSequence):
+ """
+ Add the request to waiting list.
+ """
+ assert not self._find_sequence(seq.request_id), f"Sequence {seq.request_id} already exists."
+ self.waiting_list.append(seq)
+
+ def _find_sequence(self, request_id: int) -> DiffusionSequence:
+ """
+ Find the request by request_id.
+ """
+ for lst in enumerate(self.waiting_list + self.running_list):
+ for seq in lst:
+ if seq.request_id == request_id:
+ return seq
+ return None
+
+ def schedule(self):
+ ret = None
+ if self._has_waiting:
+ ret = self.waiting_list[0]
+ self.waiting_list = self.waiting_list[1:]
+ return ret
+
+
+class RequestHandler(NaiveRequestHandler):
"""
RequestHandler is the core for handling existing requests and updating current batch.
During generation process, we call schedule function each iteration to update current batch.
@@ -176,12 +215,12 @@ class RequestHandler:
generated_token_size=inference_config.generated_token_size,
)
+ def _has_running(self) -> bool:
+ return not self.running_bb.is_empty()
+
def _init_cache(self, model_config):
self.cache_manager = KVCacheManager(self.inference_config, model_config)
- def _has_waiting(self) -> bool:
- return any(lst for lst in self.waiting_list)
-
def get_kvcache(self):
return self.cache_manager.get_kv_cache()
@@ -318,7 +357,7 @@ class RequestHandler:
if seq.output_token_id[-1] == generation_config.eos_token_id or seq.output_len >= max_new_tokens:
seq.mark_finished()
- def check_unfinished_seqs(self) -> bool:
+ def check_unfinished_reqs(self) -> bool:
return self._has_waiting() or not self.running_list.is_empty()
def total_requests_in_batch_bucket(self) -> int:
diff --git a/colossalai/inference/modeling/layers/diffusion.py b/colossalai/inference/modeling/layers/diffusion.py
new file mode 100644
index 000000000..9dc90733d
--- /dev/null
+++ b/colossalai/inference/modeling/layers/diffusion.py
@@ -0,0 +1,54 @@
+import inspect
+import types
+
+import torch
+from torch import nn
+
+
+class DiffusionPipe(nn.Module):
+ """
+ This Class convert a class of `DiffusionPipeline` into `nn.Module` and reserve most of origin attr,function and property.
+ """
+
+ def __init__(self, source_obj) -> None:
+ super(DiffusionPipe, self).__init__()
+
+ for k, v in source_obj.__dict__.items():
+ if isinstance(v, nn.Module):
+ self.add_module(k, v)
+ else:
+ setattr(self, k, v)
+
+ skip_list = ["_execution_device", "to", "device"] # this
+
+ for name, member in inspect.getmembers(source_obj.__class__):
+ if name in skip_list:
+ continue
+ if not name.startswith("__") and not name.endswith("__"):
+ if isinstance(member, property):
+ setattr(self.__class__, name, member)
+ elif inspect.isfunction(member) or inspect.ismethod(member):
+ bound_method = types.MethodType(member, self)
+ setattr(self, name, bound_method)
+ elif not callable(member) and not isinstance(member, property):
+ setattr(self, name, member)
+ elif name == "__call__":
+ bound_method = types.MethodType(member, self)
+ setattr(self, "_forward", bound_method)
+
+ @property
+ def _execution_device(self):
+ r"""
+ Returns the device on which the pipeline's models will be executed. After calling
+ [`~DiffusionPipeline.enable_sequential_cpu_offload`] the execution device can only be inferred from
+ Accelerate's module hooks.
+ """
+ # return self.device
+ return torch.device("cuda")
+
+ @property
+ def device(self):
+ next(self.parameters()).device
+
+ def forward(self, *args, **kwargs):
+ return self._forward(*args, **kwargs)
diff --git a/colossalai/inference/modeling/layers/distrifusion.py b/colossalai/inference/modeling/layers/distrifusion.py
new file mode 100644
index 000000000..ea97cceef
--- /dev/null
+++ b/colossalai/inference/modeling/layers/distrifusion.py
@@ -0,0 +1,626 @@
+# Code refer and adapted from:
+# https://github.com/huggingface/diffusers/blob/v0.29.0-release/src/diffusers
+# https://github.com/PipeFusion/PipeFusion
+
+import inspect
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from diffusers.models import attention_processor
+from diffusers.models.attention import Attention
+from diffusers.models.embeddings import PatchEmbed, get_2d_sincos_pos_embed
+from diffusers.models.transformers.pixart_transformer_2d import PixArtTransformer2DModel
+from diffusers.models.transformers.transformer_sd3 import SD3Transformer2DModel
+from torch import nn
+from torch.distributed import ProcessGroup
+
+from colossalai.inference.config import ModelShardInferenceConfig
+from colossalai.logging import get_dist_logger
+from colossalai.shardformer.layer.parallel_module import ParallelModule
+from colossalai.utils import get_current_device
+
+try:
+ from flash_attn import flash_attn_func
+
+ HAS_FLASH_ATTN = True
+except ImportError:
+ HAS_FLASH_ATTN = False
+
+
+logger = get_dist_logger(__name__)
+
+
+# adapted from https://github.com/huggingface/diffusers/blob/v0.29.0-release/src/diffusers/models/transformers/transformer_2d.py
+def PixArtAlphaTransformer2DModel_forward(
+ self: PixArtTransformer2DModel,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ timestep: Optional[torch.LongTensor] = None,
+ added_cond_kwargs: Dict[str, torch.Tensor] = None,
+ class_labels: Optional[torch.LongTensor] = None,
+ cross_attention_kwargs: Dict[str, Any] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ encoder_attention_mask: Optional[torch.Tensor] = None,
+ return_dict: bool = True,
+):
+ assert hasattr(
+ self, "patched_parallel_size"
+ ), "please check your policy, `Transformer2DModel` Must have attribute `patched_parallel_size`"
+
+ if cross_attention_kwargs is not None:
+ if cross_attention_kwargs.get("scale", None) is not None:
+ logger.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.")
+ # ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
+ # we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
+ # we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
+ # expects mask of shape:
+ # [batch, key_tokens]
+ # adds singleton query_tokens dimension:
+ # [batch, 1, key_tokens]
+ # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
+ # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn)
+ # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
+ if attention_mask is not None and attention_mask.ndim == 2:
+ # assume that mask is expressed as:
+ # (1 = keep, 0 = discard)
+ # convert mask into a bias that can be added to attention scores:
+ # (keep = +0, discard = -10000.0)
+ attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0
+ attention_mask = attention_mask.unsqueeze(1)
+
+ # convert encoder_attention_mask to a bias the same way we do for attention_mask
+ if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
+ encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0
+ encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
+
+ # 1. Input
+ batch_size = hidden_states.shape[0]
+ height, width = (
+ hidden_states.shape[-2] // self.config.patch_size,
+ hidden_states.shape[-1] // self.config.patch_size,
+ )
+ hidden_states = self.pos_embed(hidden_states)
+
+ timestep, embedded_timestep = self.adaln_single(
+ timestep, added_cond_kwargs, batch_size=batch_size, hidden_dtype=hidden_states.dtype
+ )
+
+ if self.caption_projection is not None:
+ encoder_hidden_states = self.caption_projection(encoder_hidden_states)
+ encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.shape[-1])
+
+ # 2. Blocks
+ for block in self.transformer_blocks:
+ hidden_states = block(
+ hidden_states,
+ attention_mask=attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_mask,
+ timestep=timestep,
+ cross_attention_kwargs=cross_attention_kwargs,
+ class_labels=class_labels,
+ )
+
+ # 3. Output
+ shift, scale = (self.scale_shift_table[None] + embedded_timestep[:, None].to(self.scale_shift_table.device)).chunk(
+ 2, dim=1
+ )
+ hidden_states = self.norm_out(hidden_states)
+ # Modulation
+ hidden_states = hidden_states * (1 + scale.to(hidden_states.device)) + shift.to(hidden_states.device)
+ hidden_states = self.proj_out(hidden_states)
+ hidden_states = hidden_states.squeeze(1)
+
+ # unpatchify
+ hidden_states = hidden_states.reshape(
+ shape=(
+ -1,
+ height // self.patched_parallel_size,
+ width,
+ self.config.patch_size,
+ self.config.patch_size,
+ self.out_channels,
+ )
+ )
+ hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
+ output = hidden_states.reshape(
+ shape=(
+ -1,
+ self.out_channels,
+ height // self.patched_parallel_size * self.config.patch_size,
+ width * self.config.patch_size,
+ )
+ )
+
+ # enable Distrifusion Optimization
+ if hasattr(self, "patched_parallel_size"):
+ from torch import distributed as dist
+
+ if (getattr(self, "output_buffer", None) is None) or (self.output_buffer.shape != output.shape):
+ self.output_buffer = torch.empty_like(output)
+ if (getattr(self, "buffer_list", None) is None) or (self.buffer_list[0].shape != output.shape):
+ self.buffer_list = [torch.empty_like(output) for _ in range(self.patched_parallel_size)]
+ output = output.contiguous()
+ dist.all_gather(self.buffer_list, output, async_op=False)
+ torch.cat(self.buffer_list, dim=2, out=self.output_buffer)
+ output = self.output_buffer
+
+ return (output,)
+
+
+# adapted from https://github.com/huggingface/diffusers/blob/v0.29.0-release/src/diffusers/models/transformers/transformer_sd3.py
+def SD3Transformer2DModel_forward(
+ self: SD3Transformer2DModel,
+ hidden_states: torch.FloatTensor,
+ encoder_hidden_states: torch.FloatTensor = None,
+ pooled_projections: torch.FloatTensor = None,
+ timestep: torch.LongTensor = None,
+ joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ return_dict: bool = True,
+) -> Union[torch.FloatTensor]:
+
+ assert hasattr(
+ self, "patched_parallel_size"
+ ), "please check your policy, `Transformer2DModel` Must have attribute `patched_parallel_size`"
+
+ height, width = hidden_states.shape[-2:]
+
+ hidden_states = self.pos_embed(hidden_states) # takes care of adding positional embeddings too.
+ temb = self.time_text_embed(timestep, pooled_projections)
+ encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+ for block in self.transformer_blocks:
+ encoder_hidden_states, hidden_states = block(
+ hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, temb=temb
+ )
+
+ hidden_states = self.norm_out(hidden_states, temb)
+ hidden_states = self.proj_out(hidden_states)
+
+ # unpatchify
+ patch_size = self.config.patch_size
+ height = height // patch_size // self.patched_parallel_size
+ width = width // patch_size
+
+ hidden_states = hidden_states.reshape(
+ shape=(hidden_states.shape[0], height, width, patch_size, patch_size, self.out_channels)
+ )
+ hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
+ output = hidden_states.reshape(
+ shape=(hidden_states.shape[0], self.out_channels, height * patch_size, width * patch_size)
+ )
+
+ # enable Distrifusion Optimization
+ if hasattr(self, "patched_parallel_size"):
+ from torch import distributed as dist
+
+ if (getattr(self, "output_buffer", None) is None) or (self.output_buffer.shape != output.shape):
+ self.output_buffer = torch.empty_like(output)
+ if (getattr(self, "buffer_list", None) is None) or (self.buffer_list[0].shape != output.shape):
+ self.buffer_list = [torch.empty_like(output) for _ in range(self.patched_parallel_size)]
+ output = output.contiguous()
+ dist.all_gather(self.buffer_list, output, async_op=False)
+ torch.cat(self.buffer_list, dim=2, out=self.output_buffer)
+ output = self.output_buffer
+
+ return (output,)
+
+
+# Code adapted from: https://github.com/PipeFusion/PipeFusion/blob/main/pipefuser/modules/dit/patch_parallel/patchembed.py
+class DistrifusionPatchEmbed(ParallelModule):
+ def __init__(
+ self,
+ module: PatchEmbed,
+ process_group: Union[ProcessGroup, List[ProcessGroup]],
+ model_shard_infer_config: ModelShardInferenceConfig = None,
+ ):
+ super().__init__()
+ self.module = module
+ self.rank = dist.get_rank(group=process_group)
+ self.patched_parallelism_size = model_shard_infer_config.patched_parallelism_size
+
+ @staticmethod
+ def from_native_module(module: PatchEmbed, process_group: Union[ProcessGroup, List[ProcessGroup]], *args, **kwargs):
+ model_shard_infer_config = kwargs.get("model_shard_infer_config", None)
+ distrifusion_embed = DistrifusionPatchEmbed(
+ module, process_group, model_shard_infer_config=model_shard_infer_config
+ )
+ return distrifusion_embed
+
+ def forward(self, latent):
+ module = self.module
+ if module.pos_embed_max_size is not None:
+ height, width = latent.shape[-2:]
+ else:
+ height, width = latent.shape[-2] // module.patch_size, latent.shape[-1] // module.patch_size
+
+ latent = module.proj(latent)
+ if module.flatten:
+ latent = latent.flatten(2).transpose(1, 2) # BCHW -> BNC
+ if module.layer_norm:
+ latent = module.norm(latent)
+ if module.pos_embed is None:
+ return latent.to(latent.dtype)
+ # Interpolate or crop positional embeddings as needed
+ if module.pos_embed_max_size:
+ pos_embed = module.cropped_pos_embed(height, width)
+ else:
+ if module.height != height or module.width != width:
+ pos_embed = get_2d_sincos_pos_embed(
+ embed_dim=module.pos_embed.shape[-1],
+ grid_size=(height, width),
+ base_size=module.base_size,
+ interpolation_scale=module.interpolation_scale,
+ )
+ pos_embed = torch.from_numpy(pos_embed).float().unsqueeze(0).to(latent.device)
+ else:
+ pos_embed = module.pos_embed
+
+ b, c, h = pos_embed.shape
+ pos_embed = pos_embed.view(b, self.patched_parallelism_size, -1, h)[:, self.rank]
+
+ return (latent + pos_embed).to(latent.dtype)
+
+
+# Code adapted from: https://github.com/PipeFusion/PipeFusion/blob/main/pipefuser/modules/dit/patch_parallel/conv2d.py
+class DistrifusionConv2D(ParallelModule):
+
+ def __init__(
+ self,
+ module: nn.Conv2d,
+ process_group: Union[ProcessGroup, List[ProcessGroup]],
+ model_shard_infer_config: ModelShardInferenceConfig = None,
+ ):
+ super().__init__()
+ self.module = module
+ self.rank = dist.get_rank(group=process_group)
+ self.patched_parallelism_size = model_shard_infer_config.patched_parallelism_size
+
+ @staticmethod
+ def from_native_module(module: nn.Conv2d, process_group: Union[ProcessGroup, List[ProcessGroup]], *args, **kwargs):
+ model_shard_infer_config = kwargs.get("model_shard_infer_config", None)
+ distrifusion_conv = DistrifusionConv2D(module, process_group, model_shard_infer_config=model_shard_infer_config)
+ return distrifusion_conv
+
+ def sliced_forward(self, x: torch.Tensor) -> torch.Tensor:
+
+ b, c, h, w = x.shape
+
+ stride = self.module.stride[0]
+ padding = self.module.padding[0]
+
+ output_h = x.shape[2] // stride // self.patched_parallelism_size
+ idx = dist.get_rank()
+ h_begin = output_h * idx * stride - padding
+ h_end = output_h * (idx + 1) * stride + padding
+ final_padding = [padding, padding, 0, 0]
+ if h_begin < 0:
+ h_begin = 0
+ final_padding[2] = padding
+ if h_end > h:
+ h_end = h
+ final_padding[3] = padding
+ sliced_input = x[:, :, h_begin:h_end, :]
+ padded_input = F.pad(sliced_input, final_padding, mode="constant")
+ return F.conv2d(
+ padded_input,
+ self.module.weight,
+ self.module.bias,
+ stride=stride,
+ padding="valid",
+ )
+
+ def forward(self, input: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+ output = self.sliced_forward(input)
+ return output
+
+
+# Code adapted from: https://github.com/huggingface/diffusers/blob/v0.29.0-release/src/diffusers/models/attention_processor.py
+class DistrifusionFusedAttention(ParallelModule):
+
+ def __init__(
+ self,
+ module: attention_processor.Attention,
+ process_group: Union[ProcessGroup, List[ProcessGroup]],
+ model_shard_infer_config: ModelShardInferenceConfig = None,
+ ):
+ super().__init__()
+ self.counter = 0
+ self.module = module
+ self.buffer_list = None
+ self.kv_buffer_idx = dist.get_rank(group=process_group)
+ self.patched_parallelism_size = model_shard_infer_config.patched_parallelism_size
+ self.handle = None
+ self.process_group = process_group
+ self.warm_step = 5 # for warmup
+
+ @staticmethod
+ def from_native_module(
+ module: attention_processor.Attention, process_group: Union[ProcessGroup, List[ProcessGroup]], *args, **kwargs
+ ) -> ParallelModule:
+ model_shard_infer_config = kwargs.get("model_shard_infer_config", None)
+ return DistrifusionFusedAttention(
+ module=module,
+ process_group=process_group,
+ model_shard_infer_config=model_shard_infer_config,
+ )
+
+ def _forward(
+ self,
+ attn: Attention,
+ hidden_states: torch.FloatTensor,
+ encoder_hidden_states: torch.FloatTensor = None,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ *args,
+ **kwargs,
+ ) -> torch.FloatTensor:
+ residual = hidden_states
+
+ input_ndim = hidden_states.ndim
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+ context_input_ndim = encoder_hidden_states.ndim
+ if context_input_ndim == 4:
+ batch_size, channel, height, width = encoder_hidden_states.shape
+ encoder_hidden_states = encoder_hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ batch_size = encoder_hidden_states.shape[0]
+
+ # `sample` projections.
+ query = attn.to_q(hidden_states)
+ key = attn.to_k(hidden_states)
+ value = attn.to_v(hidden_states)
+
+ kv = torch.cat([key, value], dim=-1) # shape of kv now: (bs, seq_len // parallel_size, dim * 2)
+
+ if self.patched_parallelism_size == 1:
+ full_kv = kv
+ else:
+ if self.buffer_list is None: # buffer not created
+ full_kv = torch.cat([kv for _ in range(self.patched_parallelism_size)], dim=1)
+ elif self.counter <= self.warm_step:
+ # logger.info(f"warmup: {self.counter}")
+ dist.all_gather(
+ self.buffer_list,
+ kv,
+ group=self.process_group,
+ async_op=False,
+ )
+ full_kv = torch.cat(self.buffer_list, dim=1)
+ else:
+ # logger.info(f"use old kv to infer: {self.counter}")
+ self.buffer_list[self.kv_buffer_idx].copy_(kv)
+ full_kv = torch.cat(self.buffer_list, dim=1)
+ assert self.handle is None, "we should maintain the kv of last step"
+ self.handle = dist.all_gather(self.buffer_list, kv, group=self.process_group, async_op=True)
+
+ key, value = torch.split(full_kv, full_kv.shape[-1] // 2, dim=-1)
+
+ # `context` projections.
+ encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
+ encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
+ encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
+
+ # attention
+ query = torch.cat([query, encoder_hidden_states_query_proj], dim=1)
+ key = torch.cat([key, encoder_hidden_states_key_proj], dim=1)
+ value = torch.cat([value, encoder_hidden_states_value_proj], dim=1)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ hidden_states = hidden_states = F.scaled_dot_product_attention(
+ query, key, value, dropout_p=0.0, is_causal=False
+ ) # NOTE(@lry89757) for torch >= 2.2, flash attn has been already integrated into scaled_dot_product_attention, https://pytorch.org/blog/pytorch2-2/
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ # Split the attention outputs.
+ hidden_states, encoder_hidden_states = (
+ hidden_states[:, : residual.shape[1]],
+ hidden_states[:, residual.shape[1] :],
+ )
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+ if not attn.context_pre_only:
+ encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+ if input_ndim == 4:
+ hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+ if context_input_ndim == 4:
+ encoder_hidden_states = encoder_hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ return hidden_states, encoder_hidden_states
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ **cross_attention_kwargs,
+ ) -> torch.Tensor:
+
+ if self.handle is not None:
+ self.handle.wait()
+ self.handle = None
+
+ b, l, c = hidden_states.shape
+ kv_shape = (b, l, self.module.to_k.out_features * 2)
+ if self.patched_parallelism_size > 1 and (self.buffer_list is None or self.buffer_list[0].shape != kv_shape):
+
+ self.buffer_list = [
+ torch.empty(kv_shape, dtype=hidden_states.dtype, device=get_current_device())
+ for _ in range(self.patched_parallelism_size)
+ ]
+
+ self.counter = 0
+
+ attn_parameters = set(inspect.signature(self.module.processor.__call__).parameters.keys())
+ quiet_attn_parameters = {"ip_adapter_masks"}
+ unused_kwargs = [
+ k for k, _ in cross_attention_kwargs.items() if k not in attn_parameters and k not in quiet_attn_parameters
+ ]
+ if len(unused_kwargs) > 0:
+ logger.warning(
+ f"cross_attention_kwargs {unused_kwargs} are not expected by {self.module.processor.__class__.__name__} and will be ignored."
+ )
+ cross_attention_kwargs = {k: w for k, w in cross_attention_kwargs.items() if k in attn_parameters}
+
+ output = self._forward(
+ self.module,
+ hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ attention_mask=attention_mask,
+ **cross_attention_kwargs,
+ )
+
+ self.counter += 1
+
+ return output
+
+
+# Code adapted from: https://github.com/PipeFusion/PipeFusion/blob/main/pipefuser/modules/dit/patch_parallel/attn.py
+class DistriSelfAttention(ParallelModule):
+ def __init__(
+ self,
+ module: Attention,
+ process_group: Union[ProcessGroup, List[ProcessGroup]],
+ model_shard_infer_config: ModelShardInferenceConfig = None,
+ ):
+ super().__init__()
+ self.counter = 0
+ self.module = module
+ self.buffer_list = None
+ self.kv_buffer_idx = dist.get_rank(group=process_group)
+ self.patched_parallelism_size = model_shard_infer_config.patched_parallelism_size
+ self.handle = None
+ self.process_group = process_group
+ self.warm_step = 3 # for warmup
+
+ @staticmethod
+ def from_native_module(
+ module: Attention, process_group: Union[ProcessGroup, List[ProcessGroup]], *args, **kwargs
+ ) -> ParallelModule:
+ model_shard_infer_config = kwargs.get("model_shard_infer_config", None)
+ return DistriSelfAttention(
+ module=module,
+ process_group=process_group,
+ model_shard_infer_config=model_shard_infer_config,
+ )
+
+ def _forward(self, hidden_states: torch.FloatTensor, scale: float = 1.0):
+ attn = self.module
+ assert isinstance(attn, Attention)
+
+ residual = hidden_states
+
+ batch_size, sequence_length, _ = hidden_states.shape
+
+ query = attn.to_q(hidden_states)
+
+ encoder_hidden_states = hidden_states
+ k = self.module.to_k(encoder_hidden_states)
+ v = self.module.to_v(encoder_hidden_states)
+ kv = torch.cat([k, v], dim=-1) # shape of kv now: (bs, seq_len // parallel_size, dim * 2)
+
+ if self.patched_parallelism_size == 1:
+ full_kv = kv
+ else:
+ if self.buffer_list is None: # buffer not created
+ full_kv = torch.cat([kv for _ in range(self.patched_parallelism_size)], dim=1)
+ elif self.counter <= self.warm_step:
+ # logger.info(f"warmup: {self.counter}")
+ dist.all_gather(
+ self.buffer_list,
+ kv,
+ group=self.process_group,
+ async_op=False,
+ )
+ full_kv = torch.cat(self.buffer_list, dim=1)
+ else:
+ # logger.info(f"use old kv to infer: {self.counter}")
+ self.buffer_list[self.kv_buffer_idx].copy_(kv)
+ full_kv = torch.cat(self.buffer_list, dim=1)
+ assert self.handle is None, "we should maintain the kv of last step"
+ self.handle = dist.all_gather(self.buffer_list, kv, group=self.process_group, async_op=True)
+
+ if HAS_FLASH_ATTN:
+ # flash attn
+ key, value = torch.split(full_kv, full_kv.shape[-1] // 2, dim=-1)
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim)
+ key = key.view(batch_size, -1, attn.heads, head_dim)
+ value = value.view(batch_size, -1, attn.heads, head_dim)
+
+ hidden_states = flash_attn_func(query, key, value, dropout_p=0.0, causal=False)
+ hidden_states = hidden_states.reshape(batch_size, -1, attn.heads * head_dim).to(query.dtype)
+ else:
+ # naive attn
+ key, value = torch.split(full_kv, full_kv.shape[-1] // 2, dim=-1)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
+ # TODO: add support for attn.scale when we move to Torch 2.1
+ hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
+
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if attn.residual_connection:
+ hidden_states = hidden_states + residual
+
+ hidden_states = hidden_states / attn.rescale_output_factor
+
+ return hidden_states
+
+ def forward(
+ self,
+ hidden_states: torch.FloatTensor,
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
+ scale: float = 1.0,
+ *args,
+ **kwargs,
+ ) -> torch.FloatTensor:
+
+ # async preallocates memo buffer
+ if self.handle is not None:
+ self.handle.wait()
+ self.handle = None
+
+ b, l, c = hidden_states.shape
+ kv_shape = (b, l, self.module.to_k.out_features * 2)
+ if self.patched_parallelism_size > 1 and (self.buffer_list is None or self.buffer_list[0].shape != kv_shape):
+
+ self.buffer_list = [
+ torch.empty(kv_shape, dtype=hidden_states.dtype, device=get_current_device())
+ for _ in range(self.patched_parallelism_size)
+ ]
+
+ self.counter = 0
+
+ output = self._forward(hidden_states, scale=scale)
+
+ self.counter += 1
+ return output
diff --git a/colossalai/inference/modeling/models/pixart_alpha.py b/colossalai/inference/modeling/models/pixart_alpha.py
new file mode 100644
index 000000000..cc2bee5ef
--- /dev/null
+++ b/colossalai/inference/modeling/models/pixart_alpha.py
@@ -0,0 +1,220 @@
+# Code adapted from:
+# https://github.com/huggingface/diffusers/blob/v0.29.0-release/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py
+
+from typing import Callable, List, Optional, Union
+
+import PIL.Image
+import torch
+from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha import (
+ ASPECT_RATIO_256_BIN,
+ ASPECT_RATIO_512_BIN,
+ ASPECT_RATIO_1024_BIN,
+)
+from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps
+
+from colossalai.logging import get_dist_logger
+
+from ..layers.diffusion import DiffusionPipe
+
+logger = get_dist_logger(__name__)
+
+
+@torch.no_grad()
+def pixart_alpha_forward(
+ self: DiffusionPipe,
+ prompt: Union[str, List[str]] = None,
+ negative_prompt: str = "",
+ num_inference_steps: int = 20,
+ timesteps: List[int] = None,
+ sigmas: List[float] = None,
+ guidance_scale: float = 4.5,
+ num_images_per_prompt: Optional[int] = 1,
+ height: Optional[int] = None,
+ width: Optional[int] = None,
+ eta: float = 0.0,
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ latents: Optional[torch.Tensor] = None,
+ prompt_embeds: Optional[torch.Tensor] = None,
+ prompt_attention_mask: Optional[torch.Tensor] = None,
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
+ negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ output_type: Optional[str] = "pil",
+ return_dict: bool = True,
+ callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+ callback_steps: int = 1,
+ clean_caption: bool = True,
+ use_resolution_binning: bool = True,
+ max_sequence_length: int = 120,
+ **kwargs,
+) -> PIL.Image:
+ # 1. Check inputs. Raise error if not correct
+ height = height or self.transformer.config.sample_size * self.vae_scale_factor
+ width = width or self.transformer.config.sample_size * self.vae_scale_factor
+ if use_resolution_binning:
+ if self.transformer.config.sample_size == 128:
+ aspect_ratio_bin = ASPECT_RATIO_1024_BIN
+ elif self.transformer.config.sample_size == 64:
+ aspect_ratio_bin = ASPECT_RATIO_512_BIN
+ elif self.transformer.config.sample_size == 32:
+ aspect_ratio_bin = ASPECT_RATIO_256_BIN
+ else:
+ raise ValueError("Invalid sample size")
+ orig_height, orig_width = height, width
+ height, width = self.image_processor.classify_height_width_bin(height, width, ratios=aspect_ratio_bin)
+
+ self.check_inputs(
+ prompt,
+ height,
+ width,
+ negative_prompt,
+ callback_steps,
+ prompt_embeds,
+ negative_prompt_embeds,
+ prompt_attention_mask,
+ negative_prompt_attention_mask,
+ )
+
+ # 2. Default height and width to transformer
+ if prompt is not None and isinstance(prompt, str):
+ batch_size = 1
+ elif prompt is not None and isinstance(prompt, list):
+ batch_size = len(prompt)
+ else:
+ batch_size = prompt_embeds.shape[0]
+
+ device = self._execution_device
+
+ # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+ # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+ # corresponds to doing no classifier free guidance.
+ do_classifier_free_guidance = guidance_scale > 1.0
+
+ # 3. Encode input prompt
+ (
+ prompt_embeds,
+ prompt_attention_mask,
+ negative_prompt_embeds,
+ negative_prompt_attention_mask,
+ ) = self.encode_prompt(
+ prompt,
+ do_classifier_free_guidance,
+ negative_prompt=negative_prompt,
+ num_images_per_prompt=num_images_per_prompt,
+ device=device,
+ prompt_embeds=prompt_embeds,
+ negative_prompt_embeds=negative_prompt_embeds,
+ prompt_attention_mask=prompt_attention_mask,
+ negative_prompt_attention_mask=negative_prompt_attention_mask,
+ clean_caption=clean_caption,
+ max_sequence_length=max_sequence_length,
+ )
+ if do_classifier_free_guidance:
+ prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+ prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
+
+ # 4. Prepare timesteps
+ timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps, sigmas)
+
+ # 5. Prepare latents.
+ latent_channels = self.transformer.config.in_channels
+ latents = self.prepare_latents(
+ batch_size * num_images_per_prompt,
+ latent_channels,
+ height,
+ width,
+ prompt_embeds.dtype,
+ device,
+ generator,
+ latents,
+ )
+
+ # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+ extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+ # 6.1 Prepare micro-conditions.
+ added_cond_kwargs = {"resolution": None, "aspect_ratio": None}
+ if self.transformer.config.sample_size == 128:
+ resolution = torch.tensor([height, width]).repeat(batch_size * num_images_per_prompt, 1)
+ aspect_ratio = torch.tensor([float(height / width)]).repeat(batch_size * num_images_per_prompt, 1)
+ resolution = resolution.to(dtype=prompt_embeds.dtype, device=device)
+ aspect_ratio = aspect_ratio.to(dtype=prompt_embeds.dtype, device=device)
+
+ if do_classifier_free_guidance:
+ resolution = torch.cat([resolution, resolution], dim=0)
+ aspect_ratio = torch.cat([aspect_ratio, aspect_ratio], dim=0)
+
+ added_cond_kwargs = {"resolution": resolution, "aspect_ratio": aspect_ratio}
+
+ # 7. Denoising loop
+ num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
+ for i, t in enumerate(timesteps):
+ latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+ latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+ current_timestep = t
+ if not torch.is_tensor(current_timestep):
+ # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+ # This would be a good case for the `match` statement (Python 3.10+)
+ is_mps = latent_model_input.device.type == "mps"
+ if isinstance(current_timestep, float):
+ dtype = torch.float32 if is_mps else torch.float64
+ else:
+ dtype = torch.int32 if is_mps else torch.int64
+ current_timestep = torch.tensor([current_timestep], dtype=dtype, device=latent_model_input.device)
+ elif len(current_timestep.shape) == 0:
+ current_timestep = current_timestep[None].to(latent_model_input.device)
+ # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+ current_timestep = current_timestep.expand(latent_model_input.shape[0])
+
+ # predict noise model_output
+ noise_pred = self.transformer(
+ latent_model_input,
+ encoder_hidden_states=prompt_embeds,
+ encoder_attention_mask=prompt_attention_mask,
+ timestep=current_timestep,
+ added_cond_kwargs=added_cond_kwargs,
+ return_dict=False,
+ )[0]
+
+ # perform guidance
+ if do_classifier_free_guidance:
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+ # learned sigma
+ if self.transformer.config.out_channels // 2 == latent_channels:
+ noise_pred = noise_pred.chunk(2, dim=1)[0]
+ else:
+ noise_pred = noise_pred
+
+ # compute previous image: x_t -> x_t-1
+ if num_inference_steps == 1:
+ # For DMD one step sampling: https://arxiv.org/abs/2311.18828
+ latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).pred_original_sample
+ else:
+ latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+
+ # call the callback, if provided
+ if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+ progress_bar.update()
+ if callback is not None and i % callback_steps == 0:
+ step_idx = i // getattr(self.scheduler, "order", 1)
+ callback(step_idx, t, latents)
+
+ output_type = "pil" # TODO(@lry89757) temporarily image, please support more return output
+ if not output_type == "latent":
+ image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
+ if use_resolution_binning:
+ image = self.image_processor.resize_and_crop_tensor(image, orig_width, orig_height)
+ else:
+ image = latents
+
+ if not output_type == "latent":
+ image = self.image_processor.postprocess(image, output_type=output_type)
+
+ # Offload all models
+ # self.maybe_free_model_hooks()
+
+ return image
diff --git a/colossalai/inference/modeling/models/stablediffusion3.py b/colossalai/inference/modeling/models/stablediffusion3.py
new file mode 100644
index 000000000..b12316403
--- /dev/null
+++ b/colossalai/inference/modeling/models/stablediffusion3.py
@@ -0,0 +1,178 @@
+# This code is adapted from huggingface diffusers: https://github.com/huggingface/diffusers/blob/v0.29.0-release/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import torch
+from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3 import retrieve_timesteps
+
+from ..layers.diffusion import DiffusionPipe
+
+
+# TODO(@lry89757) temporarily image, please support more return output
+@torch.no_grad()
+def sd3_forward(
+ self: DiffusionPipe,
+ prompt: Union[str, List[str]] = None,
+ prompt_2: Optional[Union[str, List[str]]] = None,
+ prompt_3: Optional[Union[str, List[str]]] = None,
+ height: Optional[int] = None,
+ width: Optional[int] = None,
+ num_inference_steps: int = 28,
+ timesteps: List[int] = None,
+ guidance_scale: float = 7.0,
+ negative_prompt: Optional[Union[str, List[str]]] = None,
+ negative_prompt_2: Optional[Union[str, List[str]]] = None,
+ negative_prompt_3: Optional[Union[str, List[str]]] = None,
+ num_images_per_prompt: Optional[int] = 1,
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ latents: Optional[torch.FloatTensor] = None,
+ prompt_embeds: Optional[torch.FloatTensor] = None,
+ negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+ pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+ negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+ output_type: Optional[str] = "pil",
+ return_dict: bool = True,
+ joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ clip_skip: Optional[int] = None,
+ callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+ callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+):
+ height = height or self.default_sample_size * self.vae_scale_factor
+ width = width or self.default_sample_size * self.vae_scale_factor
+
+ # 1. Check inputs. Raise error if not correct
+ self.check_inputs(
+ prompt,
+ prompt_2,
+ prompt_3,
+ height,
+ width,
+ negative_prompt=negative_prompt,
+ negative_prompt_2=negative_prompt_2,
+ negative_prompt_3=negative_prompt_3,
+ prompt_embeds=prompt_embeds,
+ negative_prompt_embeds=negative_prompt_embeds,
+ pooled_prompt_embeds=pooled_prompt_embeds,
+ negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+ callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+ )
+
+ self._guidance_scale = guidance_scale
+ self._clip_skip = clip_skip
+ self._joint_attention_kwargs = joint_attention_kwargs
+ self._interrupt = False
+
+ # 2. Define call parameters
+ if prompt is not None and isinstance(prompt, str):
+ batch_size = 1
+ elif prompt is not None and isinstance(prompt, list):
+ batch_size = len(prompt)
+ else:
+ batch_size = prompt_embeds.shape[0]
+
+ device = self._execution_device
+
+ (
+ prompt_embeds,
+ negative_prompt_embeds,
+ pooled_prompt_embeds,
+ negative_pooled_prompt_embeds,
+ ) = self.encode_prompt(
+ prompt=prompt,
+ prompt_2=prompt_2,
+ prompt_3=prompt_3,
+ negative_prompt=negative_prompt,
+ negative_prompt_2=negative_prompt_2,
+ negative_prompt_3=negative_prompt_3,
+ do_classifier_free_guidance=self.do_classifier_free_guidance,
+ prompt_embeds=prompt_embeds,
+ negative_prompt_embeds=negative_prompt_embeds,
+ pooled_prompt_embeds=pooled_prompt_embeds,
+ negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+ device=device,
+ clip_skip=self.clip_skip,
+ num_images_per_prompt=num_images_per_prompt,
+ )
+
+ if self.do_classifier_free_guidance:
+ prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+ pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0)
+
+ # 4. Prepare timesteps
+ timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
+ num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+ self._num_timesteps = len(timesteps)
+
+ # 5. Prepare latent variables
+ num_channels_latents = self.transformer.config.in_channels
+ latents = self.prepare_latents(
+ batch_size * num_images_per_prompt,
+ num_channels_latents,
+ height,
+ width,
+ prompt_embeds.dtype,
+ device,
+ generator,
+ latents,
+ )
+
+ # 6. Denoising loop
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
+ for i, t in enumerate(timesteps):
+ if self.interrupt:
+ continue
+
+ # expand the latents if we are doing classifier free guidance
+ latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
+ # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+ timestep = t.expand(latent_model_input.shape[0])
+
+ noise_pred = self.transformer(
+ hidden_states=latent_model_input,
+ timestep=timestep,
+ encoder_hidden_states=prompt_embeds,
+ pooled_projections=pooled_prompt_embeds,
+ joint_attention_kwargs=self.joint_attention_kwargs,
+ return_dict=False,
+ )[0]
+
+ # perform guidance
+ if self.do_classifier_free_guidance:
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+ noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+ # compute the previous noisy sample x_t -> x_t-1
+ latents_dtype = latents.dtype
+ latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+ if latents.dtype != latents_dtype:
+ if torch.backends.mps.is_available():
+ # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+ latents = latents.to(latents_dtype)
+
+ if callback_on_step_end is not None:
+ callback_kwargs = {}
+ for k in callback_on_step_end_tensor_inputs:
+ callback_kwargs[k] = locals()[k]
+ callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+ latents = callback_outputs.pop("latents", latents)
+ prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+ negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
+ negative_pooled_prompt_embeds = callback_outputs.pop(
+ "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds
+ )
+
+ # call the callback, if provided
+ if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+ progress_bar.update()
+
+ if output_type == "latent":
+ image = latents
+
+ else:
+ latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+
+ image = self.vae.decode(latents, return_dict=False)[0]
+ image = self.image_processor.postprocess(image, output_type=output_type)
+
+ return image
diff --git a/colossalai/inference/modeling/policy/__init__.py b/colossalai/inference/modeling/policy/__init__.py
index fa0395590..02ffadd9f 100644
--- a/colossalai/inference/modeling/policy/__init__.py
+++ b/colossalai/inference/modeling/policy/__init__.py
@@ -1,16 +1,22 @@
from .glide_llama import GlideLlamaModelPolicy
from .nopadding_baichuan import NoPaddingBaichuanModelInferPolicy
from .nopadding_llama import NoPaddingLlamaModelInferPolicy
+from .pixart_alpha import PixArtAlphaInferPolicy
+from .stablediffusion3 import StableDiffusion3InferPolicy
model_policy_map = {
"nopadding_llama": NoPaddingLlamaModelInferPolicy,
"nopadding_baichuan": NoPaddingBaichuanModelInferPolicy,
"glide_llama": GlideLlamaModelPolicy,
+ "StableDiffusion3Pipeline": StableDiffusion3InferPolicy,
+ "PixArtAlphaPipeline": PixArtAlphaInferPolicy,
}
__all__ = [
"NoPaddingLlamaModelInferPolicy",
"NoPaddingBaichuanModelInferPolicy",
"GlideLlamaModelPolicy",
+ "StableDiffusion3InferPolicy",
+ "PixArtAlphaInferPolicy",
"model_polic_map",
]
diff --git a/colossalai/inference/modeling/policy/pixart_alpha.py b/colossalai/inference/modeling/policy/pixart_alpha.py
new file mode 100644
index 000000000..1150b2432
--- /dev/null
+++ b/colossalai/inference/modeling/policy/pixart_alpha.py
@@ -0,0 +1,79 @@
+from diffusers.models.attention import BasicTransformerBlock
+from diffusers.models.transformers.pixart_transformer_2d import PixArtTransformer2DModel
+from torch import nn
+
+from colossalai.inference.config import RPC_PARAM
+from colossalai.inference.modeling.layers.diffusion import DiffusionPipe
+from colossalai.inference.modeling.layers.distrifusion import (
+ DistrifusionConv2D,
+ DistrifusionPatchEmbed,
+ DistriSelfAttention,
+ PixArtAlphaTransformer2DModel_forward,
+)
+from colossalai.inference.modeling.models.pixart_alpha import pixart_alpha_forward
+from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription
+
+
+class PixArtAlphaInferPolicy(Policy, RPC_PARAM):
+ def __init__(self) -> None:
+ super().__init__()
+
+ def module_policy(self):
+ policy = {}
+
+ if self.shard_config.extra_kwargs["model_shard_infer_config"].patched_parallelism_size > 1:
+
+ policy[PixArtTransformer2DModel] = ModulePolicyDescription(
+ sub_module_replacement=[
+ SubModuleReplacementDescription(
+ suffix="pos_embed.proj",
+ target_module=DistrifusionConv2D,
+ kwargs={"model_shard_infer_config": self.shard_config.extra_kwargs["model_shard_infer_config"]},
+ ),
+ SubModuleReplacementDescription(
+ suffix="pos_embed",
+ target_module=DistrifusionPatchEmbed,
+ kwargs={"model_shard_infer_config": self.shard_config.extra_kwargs["model_shard_infer_config"]},
+ ),
+ ],
+ attribute_replacement={
+ "patched_parallel_size": self.shard_config.extra_kwargs[
+ "model_shard_infer_config"
+ ].patched_parallelism_size
+ },
+ method_replacement={"forward": PixArtAlphaTransformer2DModel_forward},
+ )
+
+ policy[BasicTransformerBlock] = ModulePolicyDescription(
+ sub_module_replacement=[
+ SubModuleReplacementDescription(
+ suffix="attn1",
+ target_module=DistriSelfAttention,
+ kwargs={
+ "model_shard_infer_config": self.shard_config.extra_kwargs["model_shard_infer_config"],
+ },
+ )
+ ]
+ )
+
+ self.append_or_create_method_replacement(
+ description={"forward": pixart_alpha_forward}, policy=policy, target_key=DiffusionPipe
+ )
+
+ return policy
+
+ def preprocess(self) -> nn.Module:
+ return self.model
+
+ def postprocess(self):
+ return self.model
+
+ def config_sanity_check(self):
+ pass
+
+ def to_rpc_param(self) -> str:
+ return __class__.__name__
+
+ @staticmethod
+ def from_rpc_param() -> "PixArtAlphaInferPolicy":
+ return PixArtAlphaInferPolicy()
diff --git a/colossalai/inference/modeling/policy/stablediffusion3.py b/colossalai/inference/modeling/policy/stablediffusion3.py
new file mode 100644
index 000000000..39b764b92
--- /dev/null
+++ b/colossalai/inference/modeling/policy/stablediffusion3.py
@@ -0,0 +1,78 @@
+from diffusers.models.attention import JointTransformerBlock
+from diffusers.models.transformers import SD3Transformer2DModel
+from torch import nn
+
+from colossalai.inference.config import RPC_PARAM
+from colossalai.inference.modeling.layers.diffusion import DiffusionPipe
+from colossalai.inference.modeling.layers.distrifusion import (
+ DistrifusionConv2D,
+ DistrifusionFusedAttention,
+ DistrifusionPatchEmbed,
+ SD3Transformer2DModel_forward,
+)
+from colossalai.inference.modeling.models.stablediffusion3 import sd3_forward
+from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription
+
+
+class StableDiffusion3InferPolicy(Policy, RPC_PARAM):
+ def __init__(self) -> None:
+ super().__init__()
+
+ def module_policy(self):
+ policy = {}
+
+ if self.shard_config.extra_kwargs["model_shard_infer_config"].patched_parallelism_size > 1:
+
+ policy[SD3Transformer2DModel] = ModulePolicyDescription(
+ sub_module_replacement=[
+ SubModuleReplacementDescription(
+ suffix="pos_embed.proj",
+ target_module=DistrifusionConv2D,
+ kwargs={"model_shard_infer_config": self.shard_config.extra_kwargs["model_shard_infer_config"]},
+ ),
+ SubModuleReplacementDescription(
+ suffix="pos_embed",
+ target_module=DistrifusionPatchEmbed,
+ kwargs={"model_shard_infer_config": self.shard_config.extra_kwargs["model_shard_infer_config"]},
+ ),
+ ],
+ attribute_replacement={
+ "patched_parallel_size": self.shard_config.extra_kwargs[
+ "model_shard_infer_config"
+ ].patched_parallelism_size
+ },
+ method_replacement={"forward": SD3Transformer2DModel_forward},
+ )
+
+ policy[JointTransformerBlock] = ModulePolicyDescription(
+ sub_module_replacement=[
+ SubModuleReplacementDescription(
+ suffix="attn",
+ target_module=DistrifusionFusedAttention,
+ kwargs={
+ "model_shard_infer_config": self.shard_config.extra_kwargs["model_shard_infer_config"],
+ },
+ )
+ ]
+ )
+
+ self.append_or_create_method_replacement(
+ description={"forward": sd3_forward}, policy=policy, target_key=DiffusionPipe
+ )
+ return policy
+
+ def preprocess(self) -> nn.Module:
+ return self.model
+
+ def postprocess(self):
+ return self.model
+
+ def config_sanity_check(self):
+ pass
+
+ def to_rpc_param(self) -> str:
+ return __class__.__name__
+
+ @staticmethod
+ def from_rpc_param() -> "StableDiffusion3InferPolicy":
+ return StableDiffusion3InferPolicy()
diff --git a/colossalai/inference/struct.py b/colossalai/inference/struct.py
index 1a3094a27..65d284296 100644
--- a/colossalai/inference/struct.py
+++ b/colossalai/inference/struct.py
@@ -2,6 +2,7 @@ import enum
from dataclasses import dataclass
from typing import Any, List
+from colossalai.inference.config import DiffusionGenerationConfig
from colossalai.logging import get_dist_logger
logger = get_dist_logger(__name__)
@@ -46,6 +47,17 @@ class RequestStatus(enum.Enum):
return status == RequestStatus.WAITING
+@dataclass
+class DiffusionSequence:
+ """
+ parameters for diffusion
+ """
+
+ request_id: int
+ prompt: str
+ generation_config: DiffusionGenerationConfig
+
+
@dataclass
class Sequence:
"""Store information of input sequence.
diff --git a/colossalai/inference/utils.py b/colossalai/inference/utils.py
index 332e84d37..d0851e362 100644
--- a/colossalai/inference/utils.py
+++ b/colossalai/inference/utils.py
@@ -5,10 +5,12 @@ Utils for model inference
import math
import os
import re
+from enum import Enum
from pathlib import Path
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Union
import torch
+from diffusers import DiffusionPipeline
from torch import nn
from colossalai.logging import get_dist_logger
@@ -159,3 +161,36 @@ def can_use_flash_attn2(dtype: torch.dtype) -> bool:
except ImportError:
logger.warning(f"flash_attn2 has not been installed yet, we will use triton flash attn instead.")
return False
+
+
+class ModelType(Enum):
+ DIFFUSION_MODEL = "Diffusion Model"
+ LLM = "Large Language Model (LLM)"
+ UNKNOWN = "Unknown Model Type"
+
+
+def get_model_type(model_or_path: Union[nn.Module, str, DiffusionPipeline]):
+ if isinstance(model_or_path, DiffusionPipeline):
+ return ModelType.DIFFUSION_MODEL
+ elif isinstance(model_or_path, nn.Module):
+ return ModelType.LLM
+ elif isinstance(model_or_path, str):
+ try:
+ from transformers import AutoConfig
+
+ hf_config = AutoConfig.from_pretrained(model_or_path, trust_remote_code=True)
+ return ModelType.LLM
+ except:
+ """
+ model type is not `ModelType.LLM`
+ """
+
+ try:
+ DiffusionPipeline.load_config(model_or_path)
+ return ModelType.DIFFUSION_MODEL
+ except:
+ """
+ model type is not `ModelType.DIFFUSION_MODEL`
+ """
+ else:
+ return ModelType.UNKNOWN
diff --git a/colossalai/initialize.py b/colossalai/initialize.py
index 71d42312e..4e2eff7ce 100644
--- a/colossalai/initialize.py
+++ b/colossalai/initialize.py
@@ -3,6 +3,12 @@
import os
+# set CUDA_DEVICE_MAX_CONNECTIONS=1 to ensure that when overlapping communication and computation,
+# the order of of kernel launches on GPUs are the same as on the CPU so that comm is launched first.
+# see https://github.com/NVIDIA/Megatron-LM/issues/533
+# https://forums.developer.nvidia.com/t/how-many-streams-maximum-number-of-streams/6571/16
+os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "1"
+
import torch.distributed as dist
from colossalai.accelerator import get_accelerator
diff --git a/colossalai/shardformer/layer/moe/__init__.py b/colossalai/legacy/moe/layer/__init__.py
similarity index 100%
rename from colossalai/shardformer/layer/moe/__init__.py
rename to colossalai/legacy/moe/layer/__init__.py
diff --git a/colossalai/shardformer/layer/moe/experts.py b/colossalai/legacy/moe/layer/experts.py
similarity index 95%
rename from colossalai/shardformer/layer/moe/experts.py
rename to colossalai/legacy/moe/layer/experts.py
index 1be7a2754..8088cf44e 100644
--- a/colossalai/shardformer/layer/moe/experts.py
+++ b/colossalai/legacy/moe/layer/experts.py
@@ -5,9 +5,9 @@ import torch
import torch.nn as nn
from colossalai.kernel.triton.llama_act_combine_kernel import HAS_TRITON
-from colossalai.moe._operation import MoeInGradScaler, MoeOutGradScaler
-from colossalai.moe.manager import MOE_MANAGER
-from colossalai.moe.utils import get_activation
+from colossalai.legacy.moe.manager import MOE_MANAGER
+from colossalai.legacy.moe.utils import get_activation
+from colossalai.moe._operation import EPGradScalerIn, EPGradScalerOut
from colossalai.shardformer.layer.utils import Randomizer
from colossalai.tensor.moe_tensor.api import get_ep_rank, get_ep_size
@@ -118,7 +118,7 @@ class MLPExperts(nn.Module):
Returns:
torch.Tensor: The output tensor of shape (num_groups, num_experts, capacity, hidden_size)
"""
- x = MoeInGradScaler.apply(x, self.ep_size)
+ x = EPGradScalerIn.apply(x, self.ep_size)
e = x.size(1)
h = x.size(-1)
@@ -157,5 +157,5 @@ class MLPExperts(nn.Module):
x = torch.cat([x[i].unsqueeze(0) for i in range(e)], dim=0)
x = x.reshape(inshape)
x = x.transpose(0, 1).contiguous()
- x = MoeOutGradScaler.apply(x, self.ep_size)
+ x = EPGradScalerOut.apply(x, self.ep_size)
return x
diff --git a/colossalai/shardformer/layer/moe/layers.py b/colossalai/legacy/moe/layer/layers.py
similarity index 99%
rename from colossalai/shardformer/layer/moe/layers.py
rename to colossalai/legacy/moe/layer/layers.py
index e5b0ef97f..e43966f68 100644
--- a/colossalai/shardformer/layer/moe/layers.py
+++ b/colossalai/legacy/moe/layer/layers.py
@@ -7,9 +7,9 @@ import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F
+from colossalai.legacy.moe.load_balance import LoadBalancer
+from colossalai.legacy.moe.utils import create_ep_hierarchical_group, get_noise_generator
from colossalai.moe._operation import AllGather, AllToAll, HierarchicalAllToAll, MoeCombine, MoeDispatch, ReduceScatter
-from colossalai.moe.load_balance import LoadBalancer
-from colossalai.moe.utils import create_ep_hierarchical_group, get_noise_generator
from colossalai.shardformer.layer.moe import MLPExperts
from colossalai.tensor.moe_tensor.api import get_dp_group, get_ep_group, get_ep_group_ranks, get_ep_size
diff --git a/colossalai/shardformer/layer/moe/routers.py b/colossalai/legacy/moe/layer/routers.py
similarity index 95%
rename from colossalai/shardformer/layer/moe/routers.py
rename to colossalai/legacy/moe/layer/routers.py
index 1be7a2754..8088cf44e 100644
--- a/colossalai/shardformer/layer/moe/routers.py
+++ b/colossalai/legacy/moe/layer/routers.py
@@ -5,9 +5,9 @@ import torch
import torch.nn as nn
from colossalai.kernel.triton.llama_act_combine_kernel import HAS_TRITON
-from colossalai.moe._operation import MoeInGradScaler, MoeOutGradScaler
-from colossalai.moe.manager import MOE_MANAGER
-from colossalai.moe.utils import get_activation
+from colossalai.legacy.moe.manager import MOE_MANAGER
+from colossalai.legacy.moe.utils import get_activation
+from colossalai.moe._operation import EPGradScalerIn, EPGradScalerOut
from colossalai.shardformer.layer.utils import Randomizer
from colossalai.tensor.moe_tensor.api import get_ep_rank, get_ep_size
@@ -118,7 +118,7 @@ class MLPExperts(nn.Module):
Returns:
torch.Tensor: The output tensor of shape (num_groups, num_experts, capacity, hidden_size)
"""
- x = MoeInGradScaler.apply(x, self.ep_size)
+ x = EPGradScalerIn.apply(x, self.ep_size)
e = x.size(1)
h = x.size(-1)
@@ -157,5 +157,5 @@ class MLPExperts(nn.Module):
x = torch.cat([x[i].unsqueeze(0) for i in range(e)], dim=0)
x = x.reshape(inshape)
x = x.transpose(0, 1).contiguous()
- x = MoeOutGradScaler.apply(x, self.ep_size)
+ x = EPGradScalerOut.apply(x, self.ep_size)
return x
diff --git a/colossalai/moe/load_balance.py b/colossalai/legacy/moe/load_balance.py
similarity index 99%
rename from colossalai/moe/load_balance.py
rename to colossalai/legacy/moe/load_balance.py
index 3dc6c02c7..7339b1a7b 100644
--- a/colossalai/moe/load_balance.py
+++ b/colossalai/legacy/moe/load_balance.py
@@ -7,7 +7,7 @@ from torch import Tensor, nn
from torch.distributed import ProcessGroup
from colossalai.cluster import ProcessGroupMesh
-from colossalai.moe.manager import MOE_MANAGER
+from colossalai.legacy.moe.manager import MOE_MANAGER
from colossalai.shardformer.layer.moe import MLPExperts
from colossalai.zero.low_level import LowLevelZeroOptimizer
diff --git a/colossalai/moe/manager.py b/colossalai/legacy/moe/manager.py
similarity index 100%
rename from colossalai/moe/manager.py
rename to colossalai/legacy/moe/manager.py
diff --git a/examples/language/openmoe/README.md b/colossalai/legacy/moe/openmoe/README.md
similarity index 100%
rename from examples/language/openmoe/README.md
rename to colossalai/legacy/moe/openmoe/README.md
diff --git a/examples/language/openmoe/benchmark/benchmark_cai.py b/colossalai/legacy/moe/openmoe/benchmark/benchmark_cai.py
similarity index 99%
rename from examples/language/openmoe/benchmark/benchmark_cai.py
rename to colossalai/legacy/moe/openmoe/benchmark/benchmark_cai.py
index b9ef915c3..5f9447246 100644
--- a/examples/language/openmoe/benchmark/benchmark_cai.py
+++ b/colossalai/legacy/moe/openmoe/benchmark/benchmark_cai.py
@@ -18,9 +18,9 @@ from colossalai.accelerator import get_accelerator
from colossalai.booster import Booster
from colossalai.booster.plugin.moe_hybrid_parallel_plugin import MoeHybridParallelPlugin
from colossalai.cluster import DistCoordinator
+from colossalai.legacy.moe.manager import MOE_MANAGER
+from colossalai.legacy.moe.utils import skip_init
from colossalai.moe.layers import apply_load_balance
-from colossalai.moe.manager import MOE_MANAGER
-from colossalai.moe.utils import skip_init
from colossalai.nn.optimizer import HybridAdam
diff --git a/examples/language/openmoe/benchmark/benchmark_cai.sh b/colossalai/legacy/moe/openmoe/benchmark/benchmark_cai.sh
similarity index 100%
rename from examples/language/openmoe/benchmark/benchmark_cai.sh
rename to colossalai/legacy/moe/openmoe/benchmark/benchmark_cai.sh
diff --git a/examples/language/openmoe/benchmark/benchmark_cai_dist.sh b/colossalai/legacy/moe/openmoe/benchmark/benchmark_cai_dist.sh
similarity index 100%
rename from examples/language/openmoe/benchmark/benchmark_cai_dist.sh
rename to colossalai/legacy/moe/openmoe/benchmark/benchmark_cai_dist.sh
diff --git a/examples/language/openmoe/benchmark/benchmark_fsdp.py b/colossalai/legacy/moe/openmoe/benchmark/benchmark_fsdp.py
similarity index 98%
rename from examples/language/openmoe/benchmark/benchmark_fsdp.py
rename to colossalai/legacy/moe/openmoe/benchmark/benchmark_fsdp.py
index b00fbd001..1ae94dd90 100644
--- a/examples/language/openmoe/benchmark/benchmark_fsdp.py
+++ b/colossalai/legacy/moe/openmoe/benchmark/benchmark_fsdp.py
@@ -14,7 +14,7 @@ from torch.utils.data.distributed import DistributedSampler
from transformers.models.llama import LlamaConfig
from utils import PerformanceEvaluator, get_model_numel
-from colossalai.moe.manager import MOE_MANAGER
+from colossalai.legacy.moe.manager import MOE_MANAGER
class RandomDataset(Dataset):
diff --git a/examples/language/openmoe/benchmark/benchmark_fsdp.sh b/colossalai/legacy/moe/openmoe/benchmark/benchmark_fsdp.sh
similarity index 100%
rename from examples/language/openmoe/benchmark/benchmark_fsdp.sh
rename to colossalai/legacy/moe/openmoe/benchmark/benchmark_fsdp.sh
diff --git a/examples/language/openmoe/benchmark/hostfile.txt b/colossalai/legacy/moe/openmoe/benchmark/hostfile.txt
similarity index 100%
rename from examples/language/openmoe/benchmark/hostfile.txt
rename to colossalai/legacy/moe/openmoe/benchmark/hostfile.txt
diff --git a/examples/language/openmoe/benchmark/utils.py b/colossalai/legacy/moe/openmoe/benchmark/utils.py
similarity index 100%
rename from examples/language/openmoe/benchmark/utils.py
rename to colossalai/legacy/moe/openmoe/benchmark/utils.py
diff --git a/examples/language/openmoe/infer.py b/colossalai/legacy/moe/openmoe/infer.py
similarity index 100%
rename from examples/language/openmoe/infer.py
rename to colossalai/legacy/moe/openmoe/infer.py
diff --git a/examples/language/openmoe/infer.sh b/colossalai/legacy/moe/openmoe/infer.sh
similarity index 100%
rename from examples/language/openmoe/infer.sh
rename to colossalai/legacy/moe/openmoe/infer.sh
diff --git a/examples/language/openmoe/model/__init__.py b/colossalai/legacy/moe/openmoe/model/__init__.py
similarity index 100%
rename from examples/language/openmoe/model/__init__.py
rename to colossalai/legacy/moe/openmoe/model/__init__.py
diff --git a/examples/language/openmoe/model/convert_openmoe_ckpt.py b/colossalai/legacy/moe/openmoe/model/convert_openmoe_ckpt.py
similarity index 100%
rename from examples/language/openmoe/model/convert_openmoe_ckpt.py
rename to colossalai/legacy/moe/openmoe/model/convert_openmoe_ckpt.py
diff --git a/examples/language/openmoe/model/convert_openmoe_ckpt.sh b/colossalai/legacy/moe/openmoe/model/convert_openmoe_ckpt.sh
similarity index 100%
rename from examples/language/openmoe/model/convert_openmoe_ckpt.sh
rename to colossalai/legacy/moe/openmoe/model/convert_openmoe_ckpt.sh
diff --git a/examples/language/openmoe/model/modeling_openmoe.py b/colossalai/legacy/moe/openmoe/model/modeling_openmoe.py
similarity index 99%
rename from examples/language/openmoe/model/modeling_openmoe.py
rename to colossalai/legacy/moe/openmoe/model/modeling_openmoe.py
index 1febacd7d..5d6e91765 100644
--- a/examples/language/openmoe/model/modeling_openmoe.py
+++ b/colossalai/legacy/moe/openmoe/model/modeling_openmoe.py
@@ -50,8 +50,8 @@ try:
except:
HAS_FLASH_ATTN = False
from colossalai.kernel.triton.llama_act_combine_kernel import HAS_TRITON
-from colossalai.moe.manager import MOE_MANAGER
-from colossalai.moe.utils import get_activation, set_moe_args
+from colossalai.legacy.moe.manager import MOE_MANAGER
+from colossalai.legacy.moe.utils import get_activation, set_moe_args
from colossalai.shardformer.layer.moe import SparseMLP
if HAS_TRITON:
diff --git a/examples/language/openmoe/model/openmoe_8b_config.json b/colossalai/legacy/moe/openmoe/model/openmoe_8b_config.json
similarity index 100%
rename from examples/language/openmoe/model/openmoe_8b_config.json
rename to colossalai/legacy/moe/openmoe/model/openmoe_8b_config.json
diff --git a/examples/language/openmoe/model/openmoe_base_config.json b/colossalai/legacy/moe/openmoe/model/openmoe_base_config.json
similarity index 100%
rename from examples/language/openmoe/model/openmoe_base_config.json
rename to colossalai/legacy/moe/openmoe/model/openmoe_base_config.json
diff --git a/examples/language/openmoe/model/openmoe_policy.py b/colossalai/legacy/moe/openmoe/model/openmoe_policy.py
similarity index 99%
rename from examples/language/openmoe/model/openmoe_policy.py
rename to colossalai/legacy/moe/openmoe/model/openmoe_policy.py
index f46062128..ccd566b08 100644
--- a/examples/language/openmoe/model/openmoe_policy.py
+++ b/colossalai/legacy/moe/openmoe/model/openmoe_policy.py
@@ -9,7 +9,7 @@ from torch.nn import Module
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.utils import logging
-from colossalai.moe.manager import MOE_MANAGER
+from colossalai.legacy.moe.manager import MOE_MANAGER
from colossalai.pipeline.stage_manager import PipelineStageManager
from colossalai.shardformer.layer import FusedRMSNorm, Linear1D_Col
from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription
diff --git a/examples/language/openmoe/requirements.txt b/colossalai/legacy/moe/openmoe/requirements.txt
similarity index 100%
rename from examples/language/openmoe/requirements.txt
rename to colossalai/legacy/moe/openmoe/requirements.txt
diff --git a/examples/language/openmoe/test_ci.sh b/colossalai/legacy/moe/openmoe/test_ci.sh
similarity index 100%
rename from examples/language/openmoe/test_ci.sh
rename to colossalai/legacy/moe/openmoe/test_ci.sh
diff --git a/examples/language/openmoe/train.py b/colossalai/legacy/moe/openmoe/train.py
similarity index 99%
rename from examples/language/openmoe/train.py
rename to colossalai/legacy/moe/openmoe/train.py
index ff0e4bad6..0173f0964 100644
--- a/examples/language/openmoe/train.py
+++ b/colossalai/legacy/moe/openmoe/train.py
@@ -19,7 +19,7 @@ from colossalai.accelerator import get_accelerator
from colossalai.booster import Booster
from colossalai.booster.plugin.moe_hybrid_parallel_plugin import MoeHybridParallelPlugin
from colossalai.cluster import DistCoordinator
-from colossalai.moe.utils import skip_init
+from colossalai.legacy.moe.utils import skip_init
from colossalai.nn.optimizer import HybridAdam
from colossalai.shardformer.layer.moe import apply_load_balance
diff --git a/examples/language/openmoe/train.sh b/colossalai/legacy/moe/openmoe/train.sh
similarity index 100%
rename from examples/language/openmoe/train.sh
rename to colossalai/legacy/moe/openmoe/train.sh
diff --git a/colossalai/moe/utils.py b/colossalai/legacy/moe/utils.py
similarity index 99%
rename from colossalai/moe/utils.py
rename to colossalai/legacy/moe/utils.py
index 3d08ab7dd..d91c41363 100644
--- a/colossalai/moe/utils.py
+++ b/colossalai/legacy/moe/utils.py
@@ -9,7 +9,7 @@ import torch.nn.functional as F
from torch.distributed.distributed_c10d import get_process_group_ranks
from colossalai.accelerator import get_accelerator
-from colossalai.moe.manager import MOE_MANAGER
+from colossalai.legacy.moe.manager import MOE_MANAGER
from colossalai.tensor.moe_tensor.api import is_moe_tensor
diff --git a/colossalai/legacy/nn/layer/parallel_1d/_operation.py b/colossalai/legacy/nn/layer/parallel_1d/_operation.py
index f01da97ba..8b8f04ccf 100644
--- a/colossalai/legacy/nn/layer/parallel_1d/_operation.py
+++ b/colossalai/legacy/nn/layer/parallel_1d/_operation.py
@@ -81,7 +81,6 @@ class LinearWithAsyncCommunication(torch.autograd.Function):
handle = dist.all_reduce(grad_input, group=gpc.get_group(ctx.parallel_mode), async_op=True)
# Delay the start of weight gradient computation shortly (3us) to have
# all-reduce scheduled first and have GPU resources allocated
- _ = torch.empty(1, device=grad_output.device) + 1
grad_weight = grad_output.t().matmul(total_input)
grad_bias = grad_output.sum(dim=0) if use_bias else None
diff --git a/colossalai/moe/__init__.py b/colossalai/moe/__init__.py
index 0623d19ef..e69de29bb 100644
--- a/colossalai/moe/__init__.py
+++ b/colossalai/moe/__init__.py
@@ -1,5 +0,0 @@
-from .manager import MOE_MANAGER
-
-__all__ = [
- "MOE_MANAGER",
-]
diff --git a/colossalai/moe/_operation.py b/colossalai/moe/_operation.py
index 01c837ee3..ac422a4da 100644
--- a/colossalai/moe/_operation.py
+++ b/colossalai/moe/_operation.py
@@ -290,7 +290,7 @@ def moe_cumsum(inputs: Tensor, use_kernel: bool = False):
return torch.cumsum(inputs, dim=0) - 1
-class MoeInGradScaler(torch.autograd.Function):
+class EPGradScalerIn(torch.autograd.Function):
"""
Scale the gradient back by the number of experts
because the batch size increases in the moe stage
@@ -298,8 +298,7 @@ class MoeInGradScaler(torch.autograd.Function):
@staticmethod
def forward(ctx: Any, inputs: Tensor, ep_size: int) -> Tensor:
- if ctx is not None:
- ctx.ep_size = ep_size
+ ctx.ep_size = ep_size
return inputs
@staticmethod
@@ -311,7 +310,7 @@ class MoeInGradScaler(torch.autograd.Function):
return grad, None
-class MoeOutGradScaler(torch.autograd.Function):
+class EPGradScalerOut(torch.autograd.Function):
"""
Scale the gradient by the number of experts
because the batch size increases in the moe stage
@@ -331,6 +330,50 @@ class MoeOutGradScaler(torch.autograd.Function):
return grad, None
+class DPGradScalerIn(torch.autograd.Function):
+ """
+ Scale the gradient back by the number of experts
+ because the batch size increases in the moe stage
+ """
+
+ @staticmethod
+ def forward(ctx: Any, inputs: Tensor, moe_dp_size: int, activated_experts: int) -> Tensor:
+ assert activated_experts != 0, f"shouldn't be called when no expert is activated"
+ ctx.moe_dp_size = moe_dp_size
+ ctx.activated_experts = activated_experts
+ return inputs
+
+ @staticmethod
+ def backward(ctx: Any, *grad_outputs: Tensor) -> Tuple[Tensor, None, None]:
+ assert len(grad_outputs) == 1
+ grad = grad_outputs[0]
+ if ctx.moe_dp_size != ctx.activated_experts:
+ grad.mul_(ctx.activated_experts / ctx.moe_dp_size)
+ return grad, None, None
+
+
+class DPGradScalerOut(torch.autograd.Function):
+ """
+ Scale the gradient by the number of experts
+ because the batch size increases in the moe stage
+ """
+
+ @staticmethod
+ def forward(ctx: Any, inputs: Tensor, moe_dp_size: int, activated_experts: int) -> Tensor:
+ assert activated_experts != 0, f"shouldn't be called when no expert is activated"
+ ctx.moe_dp_size = moe_dp_size
+ ctx.activated_experts = activated_experts
+ return inputs
+
+ @staticmethod
+ def backward(ctx: Any, *grad_outputs: Tensor) -> Tuple[Tensor, None, None]:
+ assert len(grad_outputs) == 1
+ grad = grad_outputs[0]
+ if ctx.moe_dp_size != ctx.activated_experts:
+ grad.mul_(ctx.moe_dp_size / ctx.activated_experts)
+ return grad, None, None
+
+
def _all_to_all(
inputs: torch.Tensor,
input_split_sizes: Optional[List[int]] = None,
@@ -393,4 +436,7 @@ def all_to_all_uneven(
group=None,
overlap: bool = False,
):
+ assert (
+ inputs.requires_grad
+ ), "Input must require grad to assure that backward is executed, otherwise it might hang the program."
return AllToAllUneven.apply(inputs, input_split_sizes, output_split_sizes, group, overlap)
diff --git a/colossalai/pipeline/p2p.py b/colossalai/pipeline/p2p.py
index ed190eb08..b7b284213 100644
--- a/colossalai/pipeline/p2p.py
+++ b/colossalai/pipeline/p2p.py
@@ -91,7 +91,11 @@ def _broadcast_object_list(
my_rank = dist.get_rank()
# Serialize object_list elements to tensors on src rank.
if my_rank == src:
- if Version(torch.__version__) >= Version("1.13.0"):
+ if Version(torch.__version__) >= Version("2.3.0"):
+ tensor_list, size_list = zip(
+ *[c10d._object_to_tensor(obj, device=current_device, group=group) for obj in object_list]
+ )
+ elif Version(torch.__version__) >= Version("1.13.0"):
tensor_list, size_list = zip(*[c10d._object_to_tensor(obj, device=current_device) for obj in object_list])
else:
tensor_list, size_list = zip(*[c10d._object_to_tensor(obj) for obj in object_list])
@@ -276,7 +280,11 @@ def _send_recv_serialization_object(
send_object_tensor = None
send_object_size_tensor = None
if object is not None and send_dst is not None:
- if Version(torch.__version__) >= Version("1.13.0"):
+ if Version(torch.__version__) >= Version("2.3.0"):
+ send_object_tensor, send_object_size_tensor = c10d._object_to_tensor(
+ object, device=current_device, group=send_group
+ )
+ elif Version(torch.__version__) >= Version("1.13.0"):
send_object_tensor, send_object_size_tensor = c10d._object_to_tensor(object, device=current_device)
else:
send_object_tensor, send_object_size_tensor = c10d._object_to_tensor(object)
diff --git a/colossalai/shardformer/layer/__init__.py b/colossalai/shardformer/layer/__init__.py
index f17fad1b6..331e49729 100644
--- a/colossalai/shardformer/layer/__init__.py
+++ b/colossalai/shardformer/layer/__init__.py
@@ -3,7 +3,7 @@ from .attn import AttnMaskType, ColoAttention
from .dropout import DropoutForParallelInput, DropoutForReplicatedInput
from .embedding import Embedding1D, PaddingEmbedding, VocabParallelEmbedding1D
from .linear import Linear1D_Col, Linear1D_Row, PaddingLMHead, VocabParallelLMHead1D
-from .loss import cross_entropy_1d
+from .loss import cross_entropy_1d, dist_cross_entropy
from .normalization import FusedLayerNorm, FusedRMSNorm, LayerNorm, RMSNorm
from .parallel_module import ParallelModule
from .qkv_fused_linear import FusedLinear1D_Col, GPT2FusedLinearConv1D_Col, GPT2FusedLinearConv1D_Row
@@ -18,6 +18,7 @@ __all__ = [
"DropoutForParallelInput",
"DropoutForReplicatedInput",
"cross_entropy_1d",
+ "dist_cross_entropy",
"BaseLayerNorm",
"LayerNorm",
"RMSNorm",
diff --git a/colossalai/shardformer/layer/_operation.py b/colossalai/shardformer/layer/_operation.py
index a27fd35c1..feebd2d05 100644
--- a/colossalai/shardformer/layer/_operation.py
+++ b/colossalai/shardformer/layer/_operation.py
@@ -146,7 +146,7 @@ class LinearWithAsyncCommunication(torch.autograd.Function):
if use_bias:
bias.view(bias.shape)
- total_input = input
+ total_input = input.contiguous()
grad_input = grad_output.matmul(weight)
grad_output = grad_output.contiguous()
# Convert the tensor shapes to 2D for execution compatibility
diff --git a/colossalai/shardformer/layer/attn.py b/colossalai/shardformer/layer/attn.py
index 141baf3d3..5872c6485 100644
--- a/colossalai/shardformer/layer/attn.py
+++ b/colossalai/shardformer/layer/attn.py
@@ -139,12 +139,11 @@ class ColoAttention:
# no padding
assert is_causal
outputs["attention_mask_type"] = AttnMaskType.CAUSAL
- attention_mask = torch.ones(s_q, s_kv, dtype=dtype, device=device).tril(diagonal=0).expand(b, s_q, s_kv)
+ attention_mask = torch.ones(s_q, s_kv, dtype=dtype, device=device)
+ if s_q != 1:
+ attention_mask = attention_mask.tril(diagonal=0)
+ attention_mask = attention_mask.expand(b, s_q, s_kv)
else:
- assert q_padding_mask.shape == (
- b,
- s_q,
- ), f"q_padding_mask shape {q_padding_mask.shape} should be the same. ({shape_4d})"
max_seqlen_q, cu_seqlens_q, q_indices = get_pad_info(q_padding_mask)
if kv_padding_mask is None:
# self attention
@@ -156,7 +155,7 @@ class ColoAttention:
b,
s_kv,
), f"q_padding_mask shape {kv_padding_mask.shape} should be the same. ({shape_4d})"
- attention_mask = q_padding_mask[:, None, :].expand(b, s_kv, s_q).to(dtype=dtype, device=device)
+ attention_mask = kv_padding_mask[:, None, :].expand(b, s_q, s_kv).to(dtype=dtype, device=device)
outputs.update(
{
"cu_seqlens_q": cu_seqlens_q,
@@ -169,7 +168,8 @@ class ColoAttention:
)
if is_causal:
outputs["attention_mask_type"] = AttnMaskType.PADDED_CAUSAL
- attention_mask = attention_mask * attention_mask.new_ones(s_q, s_kv).tril(diagonal=0)
+ if s_q != 1:
+ attention_mask = attention_mask * attention_mask.new_ones(s_q, s_kv).tril(diagonal=0)
else:
outputs["attention_mask_type"] = AttnMaskType.PADDED
attention_mask = invert_mask(attention_mask).unsqueeze(1)
diff --git a/colossalai/shardformer/layer/loss.py b/colossalai/shardformer/layer/loss.py
index a6d19edf5..cea2da03f 100644
--- a/colossalai/shardformer/layer/loss.py
+++ b/colossalai/shardformer/layer/loss.py
@@ -2,8 +2,11 @@ import torch
import torch.distributed as dist
from torch.autograd import Function
from torch.distributed import ProcessGroup
+from torch.nn import CrossEntropyLoss
-__all__ = ["DistCrossEntropy", "cross_entropy_1d"]
+from colossalai.shardformer.shard import ShardConfig
+
+__all__ = ["DistCrossEntropy", "cross_entropy_1d", "dist_cross_entropy"]
class DistCrossEntropy(Function):
@@ -132,3 +135,43 @@ def cross_entropy_1d(
dtype: torch.dtype = None,
) -> torch.Tensor:
return DistCrossEntropy.apply(vocab_logits, labels, ignore_index, process_group, vocab_size, dtype)
+
+
+def dist_cross_entropy(
+ labels: torch.Tensor,
+ logits: torch.Tensor,
+ shard_config: ShardConfig,
+ out_features: int,
+ vocab_size: int,
+ dtype: torch.dtype,
+) -> torch.Tensor:
+ """
+ Helper to compute cross entropy loss for most shardformer models,
+ compatible with PP, TP and SP.
+ """
+ if labels is not None:
+ # Shift so that tokens < n predict n
+ shift_logits = logits[..., :-1, :].contiguous()
+ shift_labels = labels[..., 1:].contiguous()
+ # Flatten the tokens
+ loss_fct = CrossEntropyLoss()
+ shift_labels = shift_labels.view(-1)
+ shift_labels = shift_labels.to(shift_logits.device)
+ if shard_config.enable_tensor_parallelism and shard_config.parallel_output:
+ # Cross entropy with all-reduce for TP
+ new_vocab_size = logits.shape[-1]
+ shift_logits = shift_logits.view(-1, new_vocab_size)
+ loss = cross_entropy_1d(
+ shift_logits,
+ shift_labels,
+ process_group=shard_config.tensor_parallel_process_group,
+ vocab_size=out_features,
+ dtype=dtype,
+ )
+ else:
+ # NOTE if use TP and not parallel_output, the output is gathered.
+ # see VocabParallelLMHead1D
+ shift_logits = shift_logits.view(-1, vocab_size)
+ loss = loss_fct(shift_logits, shift_labels)
+
+ return loss
diff --git a/colossalai/shardformer/layer/qkv_fused_linear.py b/colossalai/shardformer/layer/qkv_fused_linear.py
index d8425b58d..93a7eb231 100644
--- a/colossalai/shardformer/layer/qkv_fused_linear.py
+++ b/colossalai/shardformer/layer/qkv_fused_linear.py
@@ -722,6 +722,7 @@ class FusedLinear1D_Col(ParallelModule):
process_group=process_group,
weight=module.weight,
bias_=module.bias,
+ n_fused=n_fused,
*args,
**kwargs,
)
diff --git a/colossalai/shardformer/modeling/bloom.py b/colossalai/shardformer/modeling/bloom.py
index 154143626..26ffef6c5 100644
--- a/colossalai/shardformer/modeling/bloom.py
+++ b/colossalai/shardformer/modeling/bloom.py
@@ -28,7 +28,7 @@ from colossalai.pipeline.stage_manager import PipelineStageManager
from colossalai.shardformer.layer._operation import gather_forward_split_backward, split_forward_gather_backward
from colossalai.shardformer.shard import ShardConfig
-from ..layer import cross_entropy_1d
+from ..layer import dist_cross_entropy
logger = logging.get_logger(__name__)
@@ -359,30 +359,14 @@ class BloomPipelineForwards:
hidden_states = transformer_outputs[0]
lm_logits = self.lm_head(hidden_states).contiguous()
- loss = None
- if labels is not None:
- # move labels to correct device to enable model parallelism
- labels = labels.to(lm_logits.device)
- # Shift so that tokens < n predict n
- shift_logits = lm_logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- batch_size, seq_length, vocab_size = shift_logits.shape
- # Flatten the tokens
- if shard_config.enable_tensor_parallelism and shard_config.parallel_output:
- new_vocab_size = lm_logits.shape[-1]
- shift_logits = shift_logits.view(-1, new_vocab_size)
- shift_labels = shift_labels.view(-1)
- loss = cross_entropy_1d(
- shift_logits,
- shift_labels,
- process_group=shard_config.tensor_parallel_process_group,
- vocab_size=self.lm_head.out_features,
- dtype=self.transformer.dtype,
- )
- else:
- loss_fct = CrossEntropyLoss()
- shift_logits = shift_logits.view(-1, self.config.vocab_size)
- loss = loss_fct(shift_logits, shift_labels.view(-1))
+ loss = dist_cross_entropy(
+ labels,
+ lm_logits,
+ shard_config,
+ self.lm_head.out_features,
+ self.config.vocab_size,
+ self.transformer.dtype,
+ )
if not return_dict:
output = (lm_logits,) + transformer_outputs[1:]
@@ -1040,24 +1024,10 @@ def get_lm_forward_with_dist_cross_entropy(shard_config: ShardConfig):
hidden_states = transformer_outputs[0]
lm_logits = self.lm_head(hidden_states)
- loss = None
- if labels is not None:
- # move labels to correct device to enable model parallelism
- labels = labels.to(lm_logits.device)
- # Shift so that tokens < n predict n
- shift_logits = lm_logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- # Flatten the tokens
- new_vocab_size = lm_logits.shape[-1]
- shift_logits = shift_logits.view(-1, new_vocab_size)
- shift_labels = shift_labels.view(-1)
- loss = cross_entropy_1d(
- shift_logits,
- shift_labels,
- process_group=shard_config.tensor_parallel_process_group,
- vocab_size=self.lm_head.out_features,
- dtype=self.transformer.dtype,
- )
+ loss = dist_cross_entropy(
+ labels, lm_logits, shard_config, self.lm_head.out_features, self.config.vocab_size, self.transformer.dtype
+ )
+
if not return_dict:
output = (lm_logits,) + transformer_outputs[1:]
return ((loss,) + output) if loss is not None else output
diff --git a/colossalai/shardformer/modeling/chatglm2.py b/colossalai/shardformer/modeling/chatglm2.py
index 53c151f02..34d900d8d 100644
--- a/colossalai/shardformer/modeling/chatglm2.py
+++ b/colossalai/shardformer/modeling/chatglm2.py
@@ -11,7 +11,11 @@ from transformers.utils import logging
from colossalai.pipeline.stage_manager import PipelineStageManager
from colossalai.shardformer import ShardConfig
from colossalai.shardformer.layer import AttnMaskType, ColoAttention
-from colossalai.shardformer.layer._operation import gather_forward_split_backward, split_forward_gather_backward
+from colossalai.shardformer.layer._operation import (
+ all_to_all_comm,
+ gather_forward_split_backward,
+ split_forward_gather_backward,
+)
def get_flash_core_attention_forward():
@@ -203,6 +207,13 @@ class ChatGLMPipelineForwards:
dim=0,
process_group=shard_config.tensor_parallel_process_group,
)
+ elif shard_config.sequence_parallelism_mode == "all_to_all":
+ hidden_states = split_forward_gather_backward(
+ hidden_states,
+ dim=0,
+ process_group=shard_config.sequence_parallel_process_group,
+ grad_scale=1 / shard_config.sequence_parallel_size,
+ )
for idx in range(start_idx, end_idx):
layer = self.encoder._get_layer(idx)
if output_hidden_states:
@@ -235,6 +246,13 @@ class ChatGLMPipelineForwards:
dim=0,
process_group=shard_config.tensor_parallel_process_group,
)
+ elif shard_config.sequence_parallelism_mode == "all_to_all":
+ hidden_states = gather_forward_split_backward(
+ hidden_states,
+ dim=0,
+ process_group=shard_config.sequence_parallel_process_group,
+ grad_scale=shard_config.sequence_parallel_size,
+ )
if output_hidden_states:
all_hidden_states = all_hidden_states + (hidden_states,)
if stage_manager.is_last_stage():
@@ -329,7 +347,9 @@ class ChatGLMPipelineForwards:
return transformer_outputs
-def get_chatglm_sequence_parallel_forward_fn(shard_config: ShardConfig):
+def get_chatglm_sequence_parallel_forward_fn(shard_config: ShardConfig, sp_mode, sp_size, sp_group):
+ logger = logging.get_logger(__name__)
+
def forward(
self,
input_ids,
@@ -381,13 +401,27 @@ def get_chatglm_sequence_parallel_forward_fn(shard_config: ShardConfig):
rotary_pos_emb = rotary_pos_emb[None, :seq_length]
rotary_pos_emb = rotary_pos_emb.transpose(0, 1).contiguous()
+ if sp_mode in ["all_to_all"] and self.training:
+ if use_cache:
+ logger.warning_once(
+ "`use_cache=True` is incompatible with sp mode `{sp_mode}`. Setting `use_cache=False`..."
+ )
+ use_cache = False
# Run encoder.
# [seq_len, batch_size, hidden_size] -> [seq_len/TP_size, batch_size, hidden_size]
- inputs_embeds = split_forward_gather_backward(
- inputs_embeds,
- dim=0,
- process_group=shard_config.tensor_parallel_process_group,
- )
+ if sp_mode in ["split_gather"]:
+ inputs_embeds = split_forward_gather_backward(
+ inputs_embeds,
+ dim=0,
+ process_group=sp_group,
+ )
+ elif sp_mode == "all_to_all":
+ inputs_embeds = split_forward_gather_backward(
+ inputs_embeds,
+ dim=0,
+ process_group=sp_group,
+ grad_scale=1 / sp_size,
+ )
hidden_states, presents, all_hidden_states, all_self_attentions = self.encoder(
inputs_embeds,
full_attention_mask,
@@ -397,11 +431,19 @@ def get_chatglm_sequence_parallel_forward_fn(shard_config: ShardConfig):
output_hidden_states=output_hidden_states,
)
- hidden_states = gather_forward_split_backward(
- hidden_states,
- dim=0,
- process_group=shard_config.tensor_parallel_process_group,
- )
+ if sp_mode in ["split_gather"]:
+ hidden_states = gather_forward_split_backward(
+ hidden_states,
+ dim=0,
+ process_group=shard_config.tensor_parallel_process_group,
+ )
+ elif sp_mode == "all_to_all":
+ hidden_states = gather_forward_split_backward(
+ hidden_states,
+ dim=0,
+ process_group=sp_group,
+ grad_scale=sp_size,
+ )
if not return_dict:
return tuple(
@@ -423,3 +465,158 @@ def get_chatglm_sequence_parallel_forward_fn(shard_config: ShardConfig):
)
return forward
+
+
+def get_chatglm_sequence_parallel_attention_forward(shard_config: ShardConfig, sp_mode, sp_size, sp_group):
+ from .chatglm2_6b.modeling_chatglm import apply_rotary_pos_emb, split_tensor_along_last_dim
+
+ def forward(
+ self,
+ hidden_states,
+ attention_mask,
+ rotary_pos_emb,
+ kv_cache=None,
+ use_cache=True,
+ ):
+ if sp_mode is not None:
+ assert sp_mode in ["all_to_all", "split_gather"], "Invalid sp_mode"
+ assert (sp_size is not None) and (
+ sp_group is not None
+ ), "Must specify sp_size and sp_group for sequence parallel"
+
+ mixed_x_layer = self.query_key_value(hidden_states)
+ if self.multi_query_attention:
+ (query_layer, key_layer, value_layer) = mixed_x_layer.split(
+ [
+ self.num_attention_heads_per_partition * self.hidden_size_per_attention_head,
+ self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,
+ self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,
+ ],
+ dim=-1,
+ )
+ query_layer = query_layer.view(
+ query_layer.size()[:-1]
+ + (
+ self.num_attention_heads_per_partition,
+ self.hidden_size_per_attention_head,
+ )
+ )
+ key_layer = key_layer.view(
+ key_layer.size()[:-1]
+ + (
+ self.num_multi_query_groups_per_partition,
+ self.hidden_size_per_attention_head,
+ )
+ )
+ value_layer = value_layer.view(
+ value_layer.size()[:-1]
+ + (
+ self.num_multi_query_groups_per_partition,
+ self.hidden_size_per_attention_head,
+ )
+ )
+ else:
+ new_tensor_shape = mixed_x_layer.size()[:-1] + (
+ self.num_attention_heads_per_partition,
+ 3 * self.hidden_size_per_attention_head,
+ )
+ mixed_x_layer = mixed_x_layer.view(*new_tensor_shape)
+ # [sq, b, np, 3 * hn] --> 3 [sq, b, np, hn]
+ (query_layer, key_layer, value_layer) = split_tensor_along_last_dim(mixed_x_layer, 3)
+
+ # sp: all-to-all comminucation when introducing sequence parallel
+ if sp_mode == "all_to_all":
+ sq, bs, _, _ = value_layer.size()
+
+ query_layer = query_layer.reshape(sq, bs, -1)
+ key_layer = key_layer.reshape(sq, bs, -1)
+ value_layer = value_layer.reshape(sq, bs, -1)
+
+ query_layer = all_to_all_comm(query_layer, sp_group, gather_dim=0)
+ key_layer = all_to_all_comm(key_layer, sp_group, gather_dim=0)
+ value_layer = all_to_all_comm(value_layer, sp_group, gather_dim=0)
+
+ query_layer = query_layer.view(
+ sq * sp_size,
+ bs,
+ self.num_attention_heads_per_partition // sp_size,
+ self.hidden_size_per_attention_head,
+ ).contiguous()
+
+ key_layer = key_layer.view(
+ sq * sp_size,
+ bs,
+ self.num_attention_heads_per_partition // sp_size,
+ self.hidden_size_per_attention_head,
+ ).contiguous()
+
+ value_layer = value_layer.view(
+ sq * sp_size,
+ bs,
+ self.num_attention_heads_per_partition // sp_size,
+ self.hidden_size_per_attention_head,
+ ).contiguous()
+
+ # apply relative positional encoding (rotary embedding)
+ if rotary_pos_emb is not None:
+ query_layer = apply_rotary_pos_emb(query_layer, rotary_pos_emb)
+ key_layer = apply_rotary_pos_emb(key_layer, rotary_pos_emb)
+
+ # adjust key and value for inference
+ if kv_cache is not None:
+ cache_k, cache_v = kv_cache
+ key_layer = torch.cat((cache_k, key_layer), dim=0)
+ value_layer = torch.cat((cache_v, value_layer), dim=0)
+ if use_cache:
+ kv_cache = (key_layer, value_layer)
+ else:
+ kv_cache = None
+
+ if self.multi_query_attention:
+ key_layer = key_layer.unsqueeze(-2)
+ key_layer = key_layer.expand(
+ -1,
+ -1,
+ -1,
+ self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition,
+ -1,
+ )
+ key_layer = key_layer.contiguous().view(
+ key_layer.size()[:2]
+ + (
+ self.num_attention_heads_per_partition,
+ self.hidden_size_per_attention_head,
+ )
+ )
+ value_layer = value_layer.unsqueeze(-2)
+ value_layer = value_layer.expand(
+ -1,
+ -1,
+ -1,
+ self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition,
+ -1,
+ )
+ value_layer = value_layer.contiguous().view(
+ value_layer.size()[:2]
+ + (
+ self.num_attention_heads_per_partition // sp_size,
+ self.hidden_size_per_attention_head,
+ )
+ )
+
+ # ==================================
+ # core attention computation
+ # ==================================
+
+ context_layer = self.core_attention(query_layer, key_layer, value_layer, attention_mask)
+ if sp_mode == "all_to_all":
+ context_layer = all_to_all_comm(context_layer, sp_group, gather_dim=2, scatter_dim=0)
+
+ # =================
+ # Output. [sq, b, h]
+ # =================
+ output = self.dense(context_layer)
+
+ return output, kv_cache
+
+ return forward
diff --git a/colossalai/shardformer/modeling/command.py b/colossalai/shardformer/modeling/command.py
index 07a7f6cbf..5b36fc7db 100644
--- a/colossalai/shardformer/modeling/command.py
+++ b/colossalai/shardformer/modeling/command.py
@@ -5,7 +5,6 @@ from typing import List, Optional, Tuple, Union
import torch
import torch.utils.checkpoint
from torch import nn
-from torch.nn import CrossEntropyLoss
from transformers.cache_utils import Cache, DynamicCache
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
from transformers.models.cohere.modeling_cohere import (
@@ -25,7 +24,7 @@ from colossalai.shardformer.layer._operation import (
)
from colossalai.shardformer.shard import ShardConfig
-from ..layer import ColoAttention, cross_entropy_1d
+from ..layer import ColoAttention, dist_cross_entropy
class CommandPipelineForwards:
@@ -117,7 +116,7 @@ class CommandPipelineForwards:
# for the other stages, hidden_states is the output of the previous stage
if shard_config.enable_flash_attention:
# in this case, attention_mask is a dict rather than a tensor
- mask_shape = (batch_size, 1, seq_length_with_past, seq_length_with_past)
+ mask_shape = (batch_size, 1, seq_length, seq_length_with_past)
attention_mask = ColoAttention.prepare_attn_kwargs(
mask_shape,
hidden_states.dtype,
@@ -135,6 +134,21 @@ class CommandPipelineForwards:
)
use_cache = False
+ if shard_config and shard_config.enable_sequence_parallelism:
+ if shard_config.sequence_parallelism_mode in ["split_gather", "ring"]:
+ hidden_states = split_forward_gather_backward(
+ hidden_states,
+ dim=1,
+ process_group=shard_config.tensor_parallel_process_group,
+ )
+ elif shard_config.sequence_parallelism_mode == "all_to_all":
+ hidden_states = split_forward_gather_backward(
+ hidden_states,
+ dim=1,
+ process_group=shard_config.sequence_parallel_process_group,
+ grad_scale=1 / shard_config.sequence_parallel_size,
+ )
+
# decoder layers
all_hidden_states = () if output_hidden_states else None
all_self_attns = () if output_attentions else None
@@ -191,6 +205,21 @@ class CommandPipelineForwards:
if stage_manager.is_last_stage():
hidden_states = self.norm(hidden_states)
+ if shard_config and shard_config.enable_sequence_parallelism:
+ if shard_config.sequence_parallelism_mode in ["split_gather", "ring"]:
+ hidden_states = gather_forward_split_backward(
+ hidden_states,
+ dim=1,
+ process_group=shard_config.tensor_parallel_process_group,
+ )
+ elif shard_config.sequence_parallelism_mode == "all_to_all":
+ hidden_states = gather_forward_split_backward(
+ hidden_states,
+ dim=1,
+ process_group=shard_config.sequence_parallel_process_group,
+ grad_scale=shard_config.sequence_parallel_size,
+ )
+
# add hidden states from the last decoder layer
if output_hidden_states:
all_hidden_states += (hidden_states,)
@@ -300,29 +329,9 @@ class CommandPipelineForwards:
logits = self.lm_head(hidden_states)
logits = logits * self.logit_scale
logits = logits.float()
- loss = None
- if labels is not None:
- # Shift so that tokens < n predict n
- shift_logits = logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- # Flatten the tokens
- loss_fct = CrossEntropyLoss()
- shift_labels = shift_labels.view(-1)
- # Enable model parallelism
- shift_labels = shift_labels.to(shift_logits.device)
- if shard_config.enable_tensor_parallelism and shard_config.parallel_output:
- new_vocab_size = logits.shape[-1]
- shift_logits = shift_logits.view(-1, new_vocab_size)
- loss = cross_entropy_1d(
- shift_logits,
- shift_labels,
- process_group=shard_config.tensor_parallel_process_group,
- vocab_size=self.lm_head.out_features,
- dtype=self.model.dtype,
- )
- else:
- shift_logits = shift_logits.view(-1, self.config.vocab_size)
- loss = loss_fct(shift_logits, shift_labels)
+ loss = dist_cross_entropy(
+ labels, logits, shard_config, self.lm_head.out_features, self.config.vocab_size, self.model.dtype
+ )
if not return_dict:
output = (logits,) + outputs[1:]
@@ -658,24 +667,14 @@ def get_lm_forward_with_dist_cross_entropy(shard_config: ShardConfig):
logits = self.lm_head(hidden_states)
logits = logits * self.logit_scale
logits = logits.float()
-
- loss = None
- if labels is not None:
- # Shift so that tokens < n predict n
- shift_logits = logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- shift_labels = shift_labels.view(-1)
- # Enable model parallelism
- shift_labels = shift_labels.to(shift_logits.device)
- new_vocab_size = logits.shape[-1]
- shift_logits = shift_logits.view(-1, new_vocab_size)
- loss = cross_entropy_1d(
- shift_logits,
- shift_labels,
- process_group=shard_config.tensor_parallel_process_group,
- vocab_size=self.lm_head.out_features,
- dtype=self.model.dtype,
- )
+ loss = dist_cross_entropy(
+ labels,
+ logits,
+ shard_config,
+ self.lm_head.out_features,
+ self.config.vocab_size,
+ self.model.dtype,
+ )
if not return_dict:
output = (logits,) + outputs[1:]
diff --git a/colossalai/shardformer/modeling/deepseek.py b/colossalai/shardformer/modeling/deepseek.py
new file mode 100644
index 000000000..a84a30972
--- /dev/null
+++ b/colossalai/shardformer/modeling/deepseek.py
@@ -0,0 +1,754 @@
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from torch.distributed import ProcessGroup
+from torch.nn import CrossEntropyLoss
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_attn_mask_utils import (
+ _prepare_4d_causal_attention_mask,
+ _prepare_4d_causal_attention_mask_for_sdpa,
+)
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.models.llama.modeling_llama import apply_rotary_pos_emb
+from transformers.utils import is_flash_attn_2_available, logging
+
+from colossalai.lazy import LazyInitContext
+from colossalai.moe._operation import (
+ DPGradScalerIn,
+ DPGradScalerOut,
+ EPGradScalerIn,
+ EPGradScalerOut,
+ all_to_all_uneven,
+)
+from colossalai.pipeline.stage_manager import PipelineStageManager
+from colossalai.shardformer.layer._operation import (
+ all_to_all_comm,
+ gather_forward_split_backward,
+ split_forward_gather_backward,
+)
+from colossalai.shardformer.layer.linear import Linear1D_Col, Linear1D_Row
+from colossalai.shardformer.shard import ShardConfig
+from colossalai.shardformer.shard.utils import set_tensors_to_none
+from colossalai.tensor.moe_tensor.api import set_moe_tensor_ep_group
+
+
+# copied from modeling_deepseek.py
+class AddAuxiliaryLoss(torch.autograd.Function):
+ """
+ The trick function of adding auxiliary (aux) loss,
+ which includes the gradient of the aux loss during backpropagation.
+ """
+
+ @staticmethod
+ def forward(ctx, x, loss):
+ assert loss.numel() == 1
+ ctx.dtype = loss.dtype
+ ctx.required_aux_loss = loss.requires_grad
+ return x
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ grad_loss = None
+ if ctx.required_aux_loss:
+ grad_loss = torch.ones(1, dtype=ctx.dtype, device=grad_output.device)
+ return grad_output, grad_loss
+
+
+class EPDeepseekMoE(nn.Module):
+ def __init__(self):
+ raise RuntimeError(f"Please use `from_native_module` to create an instance of {self.__class__.__name__}")
+
+ def setup_process_groups(self, tp_group: ProcessGroup, moe_dp_group: ProcessGroup, ep_group: ProcessGroup):
+ assert tp_group is not None
+ assert moe_dp_group is not None
+ assert ep_group is not None
+
+ self.ep_size = dist.get_world_size(ep_group)
+ self.ep_rank = dist.get_rank(ep_group)
+ self.num_experts = self.config.n_routed_experts
+ assert self.num_experts % self.ep_size == 0
+
+ self.ep_group = ep_group
+ self.num_experts_per_ep = self.num_experts // self.ep_size
+ self.expert_start_idx = self.ep_rank * self.num_experts_per_ep
+ held_experts = self.experts[self.expert_start_idx : self.expert_start_idx + self.num_experts_per_ep]
+
+ set_tensors_to_none(self.experts, exclude=set(held_experts))
+
+ # setup moe_dp group
+ self.moe_dp_group = moe_dp_group
+ self.moe_dp_size = moe_dp_group.size()
+
+ # setup tp group
+ self.tp_group = tp_group
+ if self.tp_group.size() > 1:
+ for expert in held_experts:
+ expert.gate_proj = Linear1D_Col.from_native_module(expert.gate_proj, self.tp_group)
+ expert.up_proj = Linear1D_Col.from_native_module(expert.up_proj, self.tp_group)
+ expert.down_proj = Linear1D_Row.from_native_module(expert.down_proj, self.tp_group)
+
+ for p in self.experts.parameters():
+ set_moe_tensor_ep_group(p, ep_group)
+
+ @staticmethod
+ def from_native_module(
+ module,
+ tp_group: ProcessGroup,
+ moe_dp_group: ProcessGroup,
+ ep_group: ProcessGroup,
+ *args,
+ **kwargs,
+ ) -> "EPDeepseekMoE":
+ LazyInitContext.materialize(module)
+ if module.__class__.__name__ == "DeepseekMLP":
+ return module
+ module.__class__ = EPDeepseekMoE
+ module.setup_process_groups(tp_group, moe_dp_group, ep_group)
+ return module
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ identity = hidden_states
+ orig_shape = hidden_states.shape
+
+ topk_experts_idx, topk_experts_weight, aux_loss = self.gate(hidden_states)
+
+ hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) # [t0, t1, t2 ...]
+ hidden_states = hidden_states.repeat_interleave(
+ self.num_experts_per_tok, dim=0
+ ) # after repeat_interleave: [t0 t0 t1 t1 t2 t2 ... ]
+
+ flat_topk_experts_idx = topk_experts_idx.view(-1) # [e0 e1 e2 ...]
+ # The elements of flat_topk_token_idx are token ids, which are arranged in ascending order of expert ids.
+ flat_topk_token_idx = flat_topk_experts_idx.argsort()
+
+ # Now we adjust the order of the hidden states, also in ascending order of expert id
+ dispatch_states = hidden_states[flat_topk_token_idx]
+ input_split_sizes = flat_topk_experts_idx.bincount(minlength=self.num_experts) # [n0, n1, n2, n3]
+ output_split_sizes = torch.zeros_like(input_split_sizes)
+
+ # [n0, n1, n2, n3] [m0, m1, m2, m3] -> [n0, n1, m0, m1] [n2, n3, m2, m3]
+ dist.all_to_all_single(output_split_sizes, input_split_sizes, group=self.ep_group)
+
+ with torch.no_grad():
+ activate_experts = output_split_sizes[: self.num_experts_per_ep].clone()
+ for i in range(1, self.ep_size):
+ activate_experts += output_split_sizes[i * self.num_experts_per_ep : (i + 1) * self.num_experts_per_ep]
+ activate_experts = (activate_experts > 0).float()
+ dist.all_reduce(activate_experts, group=self.moe_dp_group)
+
+ input_split_list = input_split_sizes.view(self.ep_size, self.num_experts_per_ep).sum(dim=-1).tolist()
+ output_split_list = output_split_sizes.view(self.ep_size, self.num_experts_per_ep).sum(dim=-1).tolist()
+ output_states, _ = all_to_all_uneven(dispatch_states, input_split_list, output_split_list, self.ep_group)
+ output_states = EPGradScalerIn.apply(output_states, self.ep_size)
+
+ if output_states.size(0) > 0:
+ if self.num_experts_per_ep == 1:
+ expert = self.experts[self.expert_start_idx]
+ output_states = DPGradScalerIn.apply(output_states, self.moe_dp_size, activate_experts[0])
+ output_states = expert(output_states)
+ output_states = DPGradScalerOut.apply(output_states, self.moe_dp_size, activate_experts[0])
+ else:
+ output_states_splits = output_states.split(output_split_sizes.tolist())
+ output_states_list = []
+ for i, split_states in enumerate(output_states_splits):
+ if split_states.size(0) == 0: # no token routed to this experts
+ continue
+ expert = self.experts[self.expert_start_idx + i % self.num_experts_per_ep]
+ split_states = DPGradScalerIn.apply(
+ split_states, self.moe_dp_size, activate_experts[i % self.num_experts_per_ep]
+ )
+ split_states = expert(split_states)
+ split_states = DPGradScalerOut.apply(
+ split_states, self.moe_dp_size, activate_experts[i % self.num_experts_per_ep]
+ )
+ output_states_list.append(split_states)
+ output_states = torch.cat(output_states_list)
+ output_states = EPGradScalerOut.apply(output_states, self.ep_size)
+ dispatch_states, _ = all_to_all_uneven(output_states, output_split_list, input_split_list, self.ep_group)
+ recover_token_idx = torch.empty_like(flat_topk_token_idx)
+ recover_token_idx[flat_topk_token_idx] = torch.arange(
+ flat_topk_token_idx.size(0), device=flat_topk_token_idx.device
+ )
+
+ output_hidden_states = dispatch_states[recover_token_idx] # t0 t0 t1 t1 t2 t2
+ output_hidden_states = output_hidden_states.view(-1, self.num_experts_per_tok, orig_shape[-1])
+ output_hidden_states = (output_hidden_states * topk_experts_weight[:, :, None]).sum(dim=-2) # (B*S, h)
+ output_hidden_states = output_hidden_states.view(*orig_shape)
+ output_hidden_states = AddAuxiliaryLoss.apply(output_hidden_states, aux_loss)
+ if self.config.n_shared_experts is not None:
+ output_hidden_states = output_hidden_states + self.shared_experts(identity)
+ return output_hidden_states
+
+
+class DeepseekPipelineForwards:
+ """
+ This class serves as a micro library for forward function substitution of Llama models
+ under pipeline setting.
+ """
+
+ @staticmethod
+ def deepseek_model_forward(
+ self: "DeepseekModel",
+ input_ids: torch.LongTensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ stage_manager: Optional[PipelineStageManager] = None,
+ hidden_states: Optional[torch.FloatTensor] = None,
+ stage_index: Optional[List[int]] = None,
+ shard_config: ShardConfig = None,
+ ):
+ r"""
+ Args:
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+ Returns:
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, AutoModelForCausalLM
+
+ >>> model = AutoModelForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
+ >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
+
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+ >>> # Generate
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+ ```"""
+ logger = logging.get_logger(__name__)
+
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+
+ output_hidden_states = (
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+ )
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ # retrieve input_ids and inputs_embeds
+ if stage_manager.is_first_stage():
+ # retrieve input_ids and inputs_embeds
+ if input_ids is not None and inputs_embeds is not None:
+ raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
+ elif input_ids is not None:
+ batch_size, seq_length = input_ids.shape
+ elif inputs_embeds is not None:
+ batch_size, seq_length, _ = inputs_embeds.shape
+ else:
+ raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+ hidden_states = inputs_embeds
+ else:
+ input_shape = hidden_states.shape[:-1]
+ batch_size, seq_length = input_shape
+ device = hidden_states.device
+
+ seq_length_with_past = seq_length
+ past_key_values_length = 0
+
+ # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future.
+ if output_attentions:
+ logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.")
+ output_attentions = False
+ if output_hidden_states:
+ logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.")
+ output_hidden_states = False
+ if use_cache:
+ logger.warning_once("use_cache=True is not supported for pipeline models at the moment.")
+ use_cache = False
+
+ if past_key_values is not None:
+ past_key_values_length = past_key_values[0][0].shape[2]
+ seq_length_with_past = seq_length_with_past + past_key_values_length
+
+ if position_ids is None:
+ position_ids = torch.arange(
+ past_key_values_length,
+ seq_length + past_key_values_length,
+ dtype=torch.long,
+ device=device,
+ )
+ position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
+ else:
+ position_ids = position_ids.view(-1, seq_length).long()
+
+ # embed positions, for the first stage, hidden_states is the input embeddings,
+ # for the other stages, hidden_states is the output of the previous stage
+ if is_flash_attn_2_available():
+ # 2d mask is passed through the layers
+ attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
+ else:
+ # 4d mask is passed through the layers
+ attention_mask = _prepare_4d_causal_attention_mask(
+ attention_mask,
+ (batch_size, seq_length),
+ hidden_states,
+ past_key_values_length,
+ sliding_window=self.config.sliding_window,
+ )
+
+ if self.gradient_checkpointing and self.training:
+ if use_cache:
+ logger.warning_once(
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+ )
+ use_cache = False
+
+ # decoder layers
+ all_hidden_states = () if output_hidden_states else None
+ all_self_attns = () if output_attentions else None
+ next_decoder_cache = None
+
+ start_idx, end_idx = stage_index[0], stage_index[1]
+ for idx, decoder_layer in enumerate(self.layers[start_idx:end_idx], start=start_idx):
+ if output_hidden_states:
+ all_hidden_states += (hidden_states,)
+
+ past_key_value = past_key_values[idx] if past_key_values is not None else None
+
+ if self.gradient_checkpointing and self.training:
+
+ def create_custom_forward(module):
+ def custom_forward(*inputs):
+ # None for past_key_value
+ return module(*inputs)
+
+ return custom_forward
+
+ layer_outputs = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(decoder_layer),
+ hidden_states,
+ attention_mask,
+ position_ids,
+ None,
+ output_attentions,
+ )
+ else:
+ layer_outputs = decoder_layer(
+ hidden_states,
+ attention_mask,
+ position_ids,
+ past_key_value,
+ output_attentions,
+ use_cache,
+ )
+
+ hidden_states = layer_outputs[0]
+
+ if use_cache:
+ next_decoder_cache = (layer_outputs[2 if output_attentions else 1],)
+ if output_attentions:
+ all_self_attns += (layer_outputs[1],)
+
+ if stage_manager.is_last_stage():
+ hidden_states = self.norm(hidden_states)
+
+ # add hidden states from the last decoder layer
+ if output_hidden_states:
+ all_hidden_states += (hidden_states,)
+ next_cache = next_decoder_cache if use_cache else None
+
+ if stage_manager.is_last_stage():
+ if not return_dict:
+ return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=next_cache,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attns,
+ )
+ # always return dict for imediate stage
+ return {
+ "hidden_states": hidden_states,
+ }
+
+ @staticmethod
+ def deepseek_for_causal_lm_forward(
+ self: "DeepseekForCausalLM",
+ input_ids: torch.LongTensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ labels: Optional[torch.LongTensor] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ stage_manager: Optional[PipelineStageManager] = None,
+ hidden_states: Optional[torch.FloatTensor] = None,
+ stage_index: Optional[List[int]] = None,
+ shard_config: ShardConfig = None,
+ ):
+ r"""
+ Args:
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+ Returns:
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, MixtralForCausalLM
+
+ >>> model = DeepseekForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
+ >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
+
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+ >>> # Generate
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+ ```"""
+ logger = logging.get_logger(__name__)
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+
+ output_hidden_states = (
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+ )
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future.
+ if output_attentions:
+ logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.")
+ output_attentions = False
+ if output_hidden_states:
+ logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.")
+ output_hidden_states = False
+
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+ outputs = DeepseekPipelineForwards.deepseek_model_forward(
+ self.model,
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ stage_manager=stage_manager,
+ hidden_states=hidden_states,
+ stage_index=stage_index,
+ )
+ past_key_values = None
+
+ if stage_manager.is_last_stage():
+ hidden_states = outputs[0]
+ logits = self.lm_head(hidden_states)
+ logits = logits.float()
+
+ loss = None
+ if labels is not None:
+ # Shift so that tokens < n predict n
+ shift_logits = logits[..., :-1, :].contiguous()
+ shift_labels = labels[..., 1:].contiguous()
+ # Flatten the tokens
+ loss_fct = CrossEntropyLoss()
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
+ shift_labels = shift_labels.view(-1)
+ # Enable model parallelism
+ shift_labels = shift_labels.to(shift_logits.device)
+ loss = loss_fct(shift_logits, shift_labels)
+
+ if not return_dict:
+ output = (logits,) + outputs[1:]
+ return (loss,) + output if loss is not None else output
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=None,
+ hidden_states=outputs[0],
+ attentions=None,
+ )
+ else:
+ out = {}
+ hidden_states = outputs.get("hidden_states")
+ out["hidden_states"] = hidden_states
+ return out
+
+
+def get_deepseek_flash_attention_forward(shard_config, sp_mode=None, sp_size=None, sp_group=None):
+ logger = logging.get_logger(__name__)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.LongTensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_value: Optional[Cache] = None,
+ output_attentions: bool = False,
+ use_cache: bool = False,
+ **kwargs,
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+ if sp_mode is not None:
+ assert sp_mode in ["all_to_all", "split_gather", "ring"], "Invalid sp_mode"
+ assert (sp_size is not None) and (
+ sp_group is not None
+ ), "Must specify sp_size and sp_group for sequence parallel"
+
+ # DeepseekFlashAttention2 attention does not support output_attentions
+ if "padding_mask" in kwargs:
+ warnings.warn(
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
+ )
+
+ # overwrite attention_mask with padding_mask
+ attention_mask = kwargs.pop("padding_mask")
+
+ output_attentions = False
+
+ bsz, q_len, _ = hidden_states.size()
+
+ # sp: modify sp_len when sequence parallel mode is ring
+ if sp_mode in ["split_gather", "ring"]:
+ q_len *= sp_size
+
+ query_states = self.q_proj(hidden_states)
+ key_states = self.k_proj(hidden_states)
+ value_states = self.v_proj(hidden_states)
+
+ # sp: all-to-all comminucation when introducing sequence parallel
+ if sp_mode == "all_to_all":
+ query_states = all_to_all_comm(query_states, sp_group)
+ key_states = all_to_all_comm(key_states, sp_group)
+ value_states = all_to_all_comm(value_states, sp_group)
+ bsz, q_len, _ = query_states.size()
+ # Flash attention requires the input to have the shape
+ # batch_size x seq_length x head_dim x hidden_dim
+ # therefore we just need to keep the original shape
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+ kv_seq_len = key_states.shape[-2]
+ if past_key_value is not None:
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+ query_states, key_states = apply_rotary_pos_emb(
+ query_states, key_states, cos, sin, position_ids, unsqueeze_dim=0
+ )
+
+ if past_key_value is not None:
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+ # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
+ # to be able to avoid many of these transpose/reshape/view.
+ query_states = query_states.transpose(1, 2)
+ key_states = key_states.transpose(1, 2)
+ value_states = value_states.transpose(1, 2)
+ dropout_rate = self.attention_dropout if self.training else 0.0
+
+ # In PEFT, usually we cast the layer norms in float32 for training stability reasons
+ # therefore the input hidden states gets silently casted in float32. Hence, we need
+ # cast them back in the correct dtype just to be sure everything works as expected.
+ # This might slowdown training & inference so it is recommended to not cast the LayerNorms
+ # in fp32. (DeepseekRMSNorm handles it correctly)
+
+ input_dtype = query_states.dtype
+ if input_dtype == torch.float32:
+ # Handle the case where the model is quantized
+ if hasattr(self.config, "_pre_quantization_dtype"):
+ target_dtype = self.config._pre_quantization_dtype
+ elif torch.is_autocast_enabled():
+ target_dtype = torch.get_autocast_gpu_dtype()
+ else:
+ target_dtype = self.q_proj.weight.dtype
+
+ logger.warning_once(
+ f"The input hidden states seems to be silently casted in float32, this might be related to"
+ f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
+ f" {target_dtype}."
+ )
+
+ query_states = query_states.to(target_dtype)
+ key_states = key_states.to(target_dtype)
+ value_states = value_states.to(target_dtype)
+ attn_output = self._flash_attention_forward(
+ query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate
+ )
+ # sp: all-to-all comminucation when introducing sequence parallel
+ if sp_mode == "all_to_all":
+ attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim).contiguous() # (1, 8, 128)
+ attn_output = all_to_all_comm(attn_output, sp_group, scatter_dim=1, gather_dim=2) # (1, 4, 256)
+ else:
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+
+ attn_output = self.o_proj(attn_output)
+
+ if not output_attentions:
+ attn_weights = None
+
+ return attn_output, attn_weights, past_key_value
+
+ return forward
+
+
+def get_deepseek_flash_attention_model_forward(shard_config, sp_mode=None, sp_size=None, sp_group=None):
+ logger = logging.get_logger(__name__)
+
+ def forward(
+ self,
+ input_ids: torch.LongTensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+ output_hidden_states = (
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+ )
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ # retrieve input_ids and inputs_embeds
+ if input_ids is not None and inputs_embeds is not None:
+ raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+ elif input_ids is not None:
+ batch_size, seq_length = input_ids.shape[:2]
+ elif inputs_embeds is not None:
+ batch_size, seq_length = inputs_embeds.shape[:2]
+ else:
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+ if self.gradient_checkpointing and self.training:
+ if use_cache:
+ logger.warning_once(
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`transformers."
+ )
+ use_cache = False
+
+ past_key_values_length = 0
+ if use_cache:
+ use_legacy_cache = not isinstance(past_key_values, Cache)
+ if use_legacy_cache:
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
+ past_key_values_length = past_key_values.get_usable_length(seq_length)
+
+ if position_ids is None:
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
+ position_ids = torch.arange(
+ past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
+ )
+ position_ids = position_ids.unsqueeze(0)
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ if self._use_flash_attention_2:
+ # 2d mask is passed through the layers
+ attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
+ elif self._use_sdpa and not output_attentions:
+ # output_attentions=True can not be supported when using SDPA, and we fall back on
+ # the manual implementation that requires a 4D causal mask in all cases.
+ attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
+ attention_mask,
+ (batch_size, seq_length),
+ inputs_embeds,
+ past_key_values_length,
+ )
+ else:
+ # 4d mask is passed through the layers
+ attention_mask = _prepare_4d_causal_attention_mask(
+ attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
+ )
+
+ if sp_mode in ["ring", "split_gather"]:
+ inputs_embeds = split_forward_gather_backward(inputs_embeds, 1, sp_group)
+ elif sp_mode == "all_to_all":
+ inputs_embeds = split_forward_gather_backward(inputs_embeds, 1, sp_group, 1 / sp_size)
+ # embed positions
+ hidden_states = inputs_embeds
+
+ # decoder layers
+ all_hidden_states = () if output_hidden_states else None
+ all_self_attns = () if output_attentions else None
+ next_decoder_cache = None
+
+ for decoder_layer in self.layers:
+ if output_hidden_states:
+ all_hidden_states += (hidden_states,)
+
+ if self.gradient_checkpointing and self.training:
+ layer_outputs = self._gradient_checkpointing_func(
+ decoder_layer.__call__,
+ hidden_states,
+ attention_mask,
+ position_ids,
+ past_key_values,
+ output_attentions,
+ use_cache,
+ )
+ else:
+ layer_outputs = decoder_layer(
+ hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_value=past_key_values,
+ output_attentions=output_attentions,
+ use_cache=use_cache,
+ )
+
+ hidden_states = layer_outputs[0]
+
+ if use_cache:
+ next_decoder_cache = layer_outputs[2 if output_attentions else 1]
+
+ if output_attentions:
+ all_self_attns += (layer_outputs[1],)
+
+ hidden_states = self.norm(hidden_states)
+
+ if sp_mode == "ring" or sp_mode == "split_gather":
+ hidden_states = gather_forward_split_backward(hidden_states, 1, sp_group)
+ elif sp_mode == "all_to_all":
+ hidden_states = gather_forward_split_backward(hidden_states, 1, sp_group, grad_scale=sp_size)
+
+ # add hidden states from the last decoder layer
+ if output_hidden_states:
+ all_hidden_states += (hidden_states,)
+
+ next_cache = None
+ if use_cache:
+ next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
+ if not return_dict:
+ return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=next_cache,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attns,
+ )
+
+ return forward
diff --git a/colossalai/shardformer/modeling/gpt2.py b/colossalai/shardformer/modeling/gpt2.py
index beaa47952..0dbf0ca5a 100644
--- a/colossalai/shardformer/modeling/gpt2.py
+++ b/colossalai/shardformer/modeling/gpt2.py
@@ -25,7 +25,7 @@ from colossalai.shardformer.layer import ColoAttention
from colossalai.shardformer.layer._operation import gather_forward_split_backward, split_forward_gather_backward
from colossalai.shardformer.shard import ShardConfig
-from ..layer import cross_entropy_1d
+from ..layer import dist_cross_entropy
logger = logging.get_logger(__name__)
@@ -372,27 +372,9 @@ class GPT2PipelineForwards:
hidden_states = outputs[0]
lm_logits = self.lm_head(hidden_states)
- loss = None
- if labels is not None:
- # move labels to correct device to enable model parallelism
- labels = labels.to(lm_logits.device)
- # Shift so that tokens < n predict n
- shift_logits = lm_logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- # Flatten the tokens
- loss_fct = CrossEntropyLoss()
- shift_logits = shift_logits.view(-1, shift_logits.size(-1))
- shift_labels = shift_labels.view(-1)
- if shard_config.enable_tensor_parallelism and shard_config.parallel_output:
- loss = cross_entropy_1d(
- shift_logits,
- shift_labels,
- process_group=shard_config.tensor_parallel_process_group,
- vocab_size=self.lm_head.out_features,
- dtype=self.transformer.dtype,
- )
- else:
- loss = loss_fct(shift_logits, shift_labels)
+ loss = dist_cross_entropy(
+ labels, lm_logits, shard_config, self.lm_head.out_features, self.config.vocab_size, self.transformer.dtype
+ )
if not return_dict:
output = (lm_logits,) + outputs[1:]
@@ -1284,24 +1266,9 @@ def get_lm_forward_with_dist_cross_entropy(shard_config: ShardConfig):
hidden_states = transformer_outputs[0]
lm_logits = self.lm_head(hidden_states)
-
- loss = None
- if labels is not None:
- # move labels to correct device to enable model parallelism
- labels = labels.to(lm_logits.device)
- # Shift so that tokens < n predict n
- shift_logits = lm_logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- # Flatten the tokens
- shift_logits = shift_logits.view(-1, shift_logits.size(-1))
- shift_labels = shift_labels.view(-1)
- loss = cross_entropy_1d(
- shift_logits,
- shift_labels,
- process_group=shard_config.tensor_parallel_process_group,
- vocab_size=self.lm_head.out_features,
- dtype=self.transformer.dtype,
- )
+ loss = dist_cross_entropy(
+ labels, lm_logits, shard_config, self.lm_head.out_features, self.config.vocab_size, self.transformer.dtype
+ )
if not return_dict:
output = (lm_logits,) + transformer_outputs[1:]
diff --git a/colossalai/shardformer/modeling/llama.py b/colossalai/shardformer/modeling/llama.py
index f83735f05..693f6584f 100644
--- a/colossalai/shardformer/modeling/llama.py
+++ b/colossalai/shardformer/modeling/llama.py
@@ -31,7 +31,7 @@ from colossalai.shardformer.layer._operation import (
)
from colossalai.shardformer.shard import ShardConfig
-from ..layer import ColoAttention, cross_entropy_1d
+from ..layer import ColoAttention, dist_cross_entropy
class LlamaPipelineForwards:
@@ -86,13 +86,20 @@ class LlamaPipelineForwards:
device = input_ids.device if input_ids is not None else inputs_embeds.device
if inputs_embeds is None:
inputs_embeds = self.embed_tokens(input_ids)
-
hidden_states = inputs_embeds
else:
input_shape = hidden_states.shape[:-1]
batch_size, seq_length = input_shape
device = hidden_states.device
+ # Support SP + PP
+ sp_mode = shard_config.sequence_parallelism_mode
+ sp_group = shard_config.sequence_parallel_process_group
+ sp_size = shard_config.sequence_parallel_size
+ if sp_mode == "all_to_all" and not stage_manager.is_first_stage():
+ # For correct positions ids. The states will be gather along the seq dim in the attention layer later.
+ seq_length *= sp_size
+
past_seen_tokens = 0
if use_cache: # kept for BC (cache positions)
if not isinstance(past_key_values, StaticCache):
@@ -101,7 +108,7 @@ class LlamaPipelineForwards:
if cache_position is None:
if isinstance(past_key_values, StaticCache):
raise ValueError("cache_position is a required argument when using StaticCache.")
- cache_position = torch.arange(past_seen_tokens, past_seen_tokens + hidden_states.shape[1], device=device)
+ cache_position = torch.arange(past_seen_tokens, past_seen_tokens + seq_length, device=device)
seq_length_with_past = seq_length + past_seen_tokens
@@ -118,7 +125,6 @@ class LlamaPipelineForwards:
if position_ids is None:
position_ids = cache_position.unsqueeze(0)
-
# embed positions, for the first stage, hidden_states is the input embeddings,
# for the other stages, hidden_states is the output of the previous stage
if shard_config.enable_flash_attention:
@@ -134,6 +140,13 @@ class LlamaPipelineForwards:
else:
attention_mask = self._update_causal_mask(attention_mask, hidden_states, cache_position)
+ # Support SP + PP
+ if stage_manager.is_first_stage():
+ if sp_mode in ["ring", "split_gather"]:
+ hidden_states = split_forward_gather_backward(hidden_states, 1, sp_group)
+ elif sp_mode == "all_to_all":
+ hidden_states = split_forward_gather_backward(hidden_states, 1, sp_group, 1 / sp_size)
+
if self.gradient_checkpointing and self.training and use_cache:
if use_cache:
logger.warning_once(
@@ -196,6 +209,10 @@ class LlamaPipelineForwards:
if stage_manager.is_last_stage():
hidden_states = self.norm(hidden_states)
+ if sp_mode == "ring" or sp_mode == "split_gather":
+ hidden_states = gather_forward_split_backward(hidden_states, 1, sp_group)
+ elif sp_mode == "all_to_all":
+ hidden_states = gather_forward_split_backward(hidden_states, 1, sp_group, grad_scale=sp_size)
# add hidden states from the last decoder layer
if output_hidden_states:
@@ -304,29 +321,9 @@ class LlamaPipelineForwards:
if stage_manager.is_last_stage():
hidden_states = outputs[0]
logits = self.lm_head(hidden_states)
- loss = None
- if labels is not None:
- # Shift so that tokens < n predict n
- shift_logits = logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- # Flatten the tokens
- loss_fct = CrossEntropyLoss()
- shift_labels = shift_labels.view(-1)
- # Enable model parallelism
- shift_labels = shift_labels.to(shift_logits.device)
- if shard_config.enable_tensor_parallelism and shard_config.parallel_output:
- new_vocab_size = logits.shape[-1]
- shift_logits = shift_logits.view(-1, new_vocab_size)
- loss = cross_entropy_1d(
- shift_logits,
- shift_labels,
- process_group=shard_config.tensor_parallel_process_group,
- vocab_size=self.lm_head.out_features,
- dtype=self.model.dtype,
- )
- else:
- shift_logits = shift_logits.view(-1, self.config.vocab_size)
- loss = loss_fct(shift_logits, shift_labels)
+ loss = dist_cross_entropy(
+ labels, logits, shard_config, self.lm_head.out_features, self.config.vocab_size, self.model.dtype
+ )
if not return_dict:
output = (logits,) + outputs[1:]
@@ -529,7 +526,6 @@ def get_llama_flash_attention_forward(shard_config: ShardConfig, sp_mode=None, s
)
kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
-
cos, sin = self.rotary_emb(value_states, position_ids)
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
@@ -649,7 +645,7 @@ def get_llama_flash_attention_model_forward(shard_config: ShardConfig, sp_mode=N
# in this case, attention_mask is a dict rather than a tensor
if shard_config.enable_flash_attention:
- mask_shape = (inputs_embeds.shape[0], 1, past_seen_tokens + seq_len, past_seen_tokens + seq_len)
+ mask_shape = (inputs_embeds.shape[0], 1, seq_len, past_seen_tokens + seq_len)
attention_mask = ColoAttention.prepare_attn_kwargs(
mask_shape,
inputs_embeds.dtype,
@@ -814,24 +810,9 @@ def get_lm_forward_with_dist_cross_entropy(shard_config: ShardConfig):
logits = self.lm_head(hidden_states)
logits = logits.float()
- loss = None
- if labels is not None:
- # Shift so that tokens < n predict n
- shift_logits = logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- shift_labels = shift_labels.view(-1)
- # Enable model parallelism
- shift_labels = shift_labels.to(shift_logits.device)
- new_vocab_size = logits.shape[-1]
- shift_logits = shift_logits.view(-1, new_vocab_size)
- loss = cross_entropy_1d(
- shift_logits,
- shift_labels,
- process_group=shard_config.tensor_parallel_process_group,
- vocab_size=self.lm_head.out_features,
- dtype=self.model.dtype,
- )
-
+ loss = dist_cross_entropy(
+ labels, logits, shard_config, self.lm_head.out_features, self.config.vocab_size, self.model.dtype
+ )
if not return_dict:
output = (logits,) + outputs[1:]
return (loss,) + output if loss is not None else output
diff --git a/colossalai/shardformer/modeling/mistral.py b/colossalai/shardformer/modeling/mistral.py
index 310c2d8e2..ec1a8a00a 100644
--- a/colossalai/shardformer/modeling/mistral.py
+++ b/colossalai/shardformer/modeling/mistral.py
@@ -19,7 +19,7 @@ from transformers.utils import logging
from colossalai.pipeline.stage_manager import PipelineStageManager
from colossalai.shardformer.shard import ShardConfig
-from ..layer import ColoAttention, cross_entropy_1d
+from ..layer import ColoAttention, dist_cross_entropy
logger = logging.get_logger(__name__)
@@ -91,7 +91,7 @@ class MistralForwards:
if shard_config.enable_flash_attention:
# in this case, attention_mask is a dict rather than a tensor
- mask_shape = (batch_size, 1, seq_length, seq_length)
+ mask_shape = (batch_size, 1, seq_length, seq_length + past_key_values_length)
attention_mask = ColoAttention.prepare_attn_kwargs(
mask_shape,
hidden_states.dtype,
@@ -275,29 +275,9 @@ class MistralForwards:
logits = self.lm_head(hidden_states)
logits = logits.float()
- loss = None
- if labels is not None:
- # Shift so that tokens < n predict n
- shift_logits = logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- # Flatten the tokens
- loss_fct = CrossEntropyLoss()
- shift_labels = shift_labels.view(-1)
- # Enable model parallelism
- shift_labels = shift_labels.to(shift_logits.device)
- if shard_config.enable_tensor_parallelism and shard_config.parallel_output:
- new_vocab_size = logits.shape[-1]
- shift_logits = shift_logits.view(-1, new_vocab_size)
- loss = cross_entropy_1d(
- shift_logits,
- shift_labels,
- process_group=shard_config.tensor_parallel_process_group,
- vocab_size=self.lm_head.out_features,
- dtype=self.model.dtype,
- )
- else:
- shift_logits = shift_logits.view(-1, self.config.vocab_size)
- loss = loss_fct(shift_logits, shift_labels)
+ loss = dist_cross_entropy(
+ labels, logits, shard_config, self.lm_head.out_features, self.config.vocab_size, self.model.dtype
+ )
if not return_dict:
output = (logits,) + outputs[1:]
@@ -708,23 +688,9 @@ def get_lm_forward_with_dist_cross_entropy(shard_config: ShardConfig):
logits = self.lm_head(hidden_states)
logits = logits.float()
- loss = None
- if labels is not None:
- # Shift so that tokens < n predict n
- shift_logits = logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- shift_labels = shift_labels.view(-1)
- # Enable model parallelism
- shift_labels = shift_labels.to(shift_logits.device)
- new_vocab_size = logits.shape[-1]
- shift_logits = shift_logits.view(-1, new_vocab_size)
- loss = cross_entropy_1d(
- shift_logits,
- shift_labels,
- process_group=shard_config.tensor_parallel_process_group,
- vocab_size=self.lm_head.out_features,
- dtype=self.model.dtype,
- )
+ loss = dist_cross_entropy(
+ labels, logits, shard_config, self.lm_head.out_features, self.config.vocab_size, self.model.dtype
+ )
if not return_dict:
output = (logits,) + outputs[1:]
diff --git a/colossalai/shardformer/modeling/mixtral.py b/colossalai/shardformer/modeling/mixtral.py
index 2fbc34302..d30ce5ea8 100644
--- a/colossalai/shardformer/modeling/mixtral.py
+++ b/colossalai/shardformer/modeling/mixtral.py
@@ -1,52 +1,105 @@
-from typing import List, Optional
+import inspect
+import warnings
+from typing import List, Optional, Tuple, Union
import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch.distributed import ProcessGroup
-
-# from colossalai.tensor.moe_tensor.moe_info import MoeParallelInfo
from torch.nn import CrossEntropyLoss
-from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_attn_mask_utils import (
+ _prepare_4d_causal_attention_mask,
+ _prepare_4d_causal_attention_mask_for_sdpa,
+)
from transformers.models.mixtral.modeling_mixtral import (
MixtralSparseMoeBlock,
MoeCausalLMOutputWithPast,
+ MoeModelOutputWithPast,
+ apply_rotary_pos_emb,
load_balancing_loss_func,
+ repeat_kv,
)
from transformers.utils import is_flash_attn_2_available, logging
from colossalai.lazy import LazyInitContext
-from colossalai.moe._operation import MoeInGradScaler, MoeOutGradScaler, all_to_all_uneven
+from colossalai.moe._operation import (
+ DPGradScalerIn,
+ DPGradScalerOut,
+ EPGradScalerIn,
+ EPGradScalerOut,
+ all_to_all_uneven,
+)
from colossalai.pipeline.stage_manager import PipelineStageManager
+from colossalai.shardformer.layer._operation import (
+ all_to_all_comm,
+ gather_forward_split_backward,
+ split_forward_gather_backward,
+)
+from colossalai.shardformer.layer.linear import Linear1D_Col, Linear1D_Row
from colossalai.shardformer.shard import ShardConfig
from colossalai.shardformer.shard.utils import set_tensors_to_none
+from colossalai.tensor.moe_tensor.api import set_moe_tensor_ep_group
+
+if is_flash_attn_2_available():
+ from flash_attn import flash_attn_func
+
+ from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
+
+ _flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters)
class EPMixtralSparseMoeBlock(MixtralSparseMoeBlock):
- def __init__(self, config):
- self.moe_info = None
- super().__init__(config)
+ def __init__(self, *args, **kwargs):
+ raise RuntimeError(f"Please use `from_native_module` to create an instance of {self.__class__.__name__}")
- def setup_ep(self, ep_group: ProcessGroup):
- ep_group = ep_group
- self.ep_size = dist.get_world_size(ep_group) if ep_group is not None else 1
- self.ep_rank = dist.get_rank(ep_group) if ep_group is not None else 0
- assert self.num_experts % self.ep_size == 0
+ def setup_process_groups(self, tp_group: ProcessGroup, moe_dp_group: ProcessGroup, ep_group: ProcessGroup):
+ assert tp_group is not None
+ assert moe_dp_group is not None
+ assert ep_group is not None
+
+ # setup ep group
+ self.ep_size = dist.get_world_size(ep_group)
+ self.ep_rank = dist.get_rank(ep_group)
self.ep_group = ep_group
+
+ if self.num_experts % self.ep_size != 0:
+ raise ValueError("The number of experts must be divisible by the number of expert parallel groups.")
+
self.num_experts_per_ep = self.num_experts // self.ep_size
self.expert_start_idx = self.ep_rank * self.num_experts_per_ep
held_experts = self.experts[self.expert_start_idx : self.expert_start_idx + self.num_experts_per_ep]
+
set_tensors_to_none(self.experts, exclude=set(held_experts))
+
+ # setup moe_dp group
+ self.moe_dp_group = moe_dp_group
+ self.moe_dp_size = moe_dp_group.size()
+
+ # setup global tp group
+ self.tp_group = tp_group
+ if self.tp_group.size() > 1:
+ for expert in held_experts:
+ expert.w1 = Linear1D_Col.from_native_module(expert.w1, self.tp_group)
+ expert.w3 = Linear1D_Col.from_native_module(expert.w3, self.tp_group)
+ expert.w2 = Linear1D_Row.from_native_module(expert.w2, self.tp_group)
+
for p in self.experts.parameters():
- p.ep_group = ep_group
+ set_moe_tensor_ep_group(p, ep_group)
@staticmethod
- def from_native_module(module: MixtralSparseMoeBlock, *args, **kwargs) -> "EPMixtralSparseMoeBlock":
+ def from_native_module(
+ module: MixtralSparseMoeBlock,
+ tp_group: ProcessGroup,
+ moe_dp_group: ProcessGroup,
+ ep_group: ProcessGroup,
+ *args,
+ **kwargs,
+ ) -> "EPMixtralSparseMoeBlock":
+ # TODO: better init
LazyInitContext.materialize(module)
module.__class__ = EPMixtralSparseMoeBlock
- # if "ep_group" in kwargs:
- assert "ep_group" in kwargs, "You should pass ep_group in SubModuleReplacementDescription via shard_config!!"
- module.setup_ep(kwargs["ep_group"])
+ module.setup_process_groups(tp_group, moe_dp_group, ep_group)
return module
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -65,20 +118,31 @@ class EPMixtralSparseMoeBlock(MixtralSparseMoeBlock):
selected_experts_idx = selected_experts.argsort()
dispatch_states = hidden_states.repeat(self.top_k, 1)[selected_experts_idx]
input_split_sizes = selected_experts.bincount(minlength=self.num_experts)
+
output_split_sizes = torch.zeros_like(input_split_sizes)
dist.all_to_all_single(output_split_sizes, input_split_sizes, group=self.ep_group)
+ with torch.no_grad():
+ activate_experts = output_split_sizes[: self.num_experts_per_ep].clone()
+ for i in range(1, self.ep_size):
+ activate_experts += output_split_sizes[i * self.num_experts_per_ep : (i + 1) * self.num_experts_per_ep]
+ activate_experts = (activate_experts > 0).float()
+ dist.all_reduce(activate_experts, group=self.moe_dp_group)
+
input_split_list = input_split_sizes.view(self.ep_size, self.num_experts_per_ep).sum(dim=-1).tolist()
output_split_list = output_split_sizes.view(self.ep_size, self.num_experts_per_ep).sum(dim=-1).tolist()
+
output_states, _ = all_to_all_uneven(dispatch_states, input_split_list, output_split_list, self.ep_group)
# compute expert output
- output_states = MoeInGradScaler.apply(output_states, self.ep_size)
+ output_states = EPGradScalerIn.apply(output_states, self.ep_size)
if output_states.size(0) > 0:
if self.num_experts_per_ep == 1:
# no need to split
expert = self.experts[self.expert_start_idx]
+ output_states = DPGradScalerIn.apply(output_states, self.moe_dp_size, activate_experts[0])
output_states = expert.act_fn(expert.w1(output_states)) * expert.w3(output_states)
output_states = expert.w2(output_states)
+ output_states = DPGradScalerOut.apply(output_states, self.moe_dp_size, activate_experts[0])
else:
output_states_splits = output_states.split(output_split_sizes.tolist())
output_states_list = []
@@ -86,12 +150,20 @@ class EPMixtralSparseMoeBlock(MixtralSparseMoeBlock):
if split_states.size(0) == 0:
continue
expert = self.experts[self.expert_start_idx + i % self.num_experts_per_ep]
+ split_states = DPGradScalerIn.apply(
+ split_states, self.moe_dp_size, activate_experts[i % self.num_experts_per_ep]
+ )
split_states = expert.act_fn(expert.w1(split_states)) * expert.w3(split_states)
split_states = expert.w2(split_states)
+ split_states = DPGradScalerOut.apply(
+ split_states, self.moe_dp_size, activate_experts[i % self.num_experts_per_ep]
+ )
output_states_list.append(split_states)
output_states = torch.cat(output_states_list)
- output_states = MoeOutGradScaler.apply(output_states, self.ep_size)
+
+ output_states = EPGradScalerOut.apply(output_states, self.ep_size)
dispatch_states, _ = all_to_all_uneven(output_states, output_split_list, input_split_list, self.ep_group)
+
recover_experts_idx = torch.empty_like(selected_experts_idx)
recover_experts_idx[selected_experts_idx] = torch.arange(
selected_experts_idx.size(0), device=selected_experts_idx.device
@@ -107,7 +179,7 @@ class EPMixtralSparseMoeBlock(MixtralSparseMoeBlock):
class MixtralPipelineForwards:
"""
- This class serves as a micro library for forward function substitution of Llama models
+ This class serves as a micro library for forward function substitution of Mixtral models
under pipeline setting.
"""
@@ -300,16 +372,29 @@ class MixtralPipelineForwards:
if output_router_logits and past_router_logits is not None:
all_router_logits = past_router_logits + all_router_logits
if stage_manager.is_last_stage():
- return tuple(
- v
- for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_logits]
- if v is not None
+ if not return_dict:
+ return tuple(
+ v
+ for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_logits]
+ if v is not None
+ )
+ return MoeModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=next_cache,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attns,
+ router_logits=all_router_logits,
)
- # always return dict for imediate stage
- return {
- "hidden_states": hidden_states,
- "past_router_logits": all_router_logits,
- }
+ else:
+ if output_router_logits:
+ return {
+ "hidden_states": hidden_states,
+ "past_router_logits": all_router_logits,
+ }
+ else:
+ return {
+ "hidden_states": hidden_states,
+ }
@staticmethod
def mixtral_for_causal_lm_forward(
@@ -441,3 +526,335 @@ class MixtralPipelineForwards:
if output_router_logits:
out["past_router_logits"] = outputs["past_router_logits"]
return out
+
+
+def get_mixtral_flash_attention_forward(shard_config, sp_mode=None, sp_size=None, sp_group=None):
+ logger = logging.get_logger(__name__)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_value: Optional[Cache] = None,
+ output_attentions: bool = False,
+ use_cache: bool = False,
+ **kwargs,
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
+ if sp_mode is not None:
+ assert sp_mode in ["all_to_all", "split_gather", "ring"], "Invalid sp_mode"
+ assert (sp_size is not None) and (
+ sp_group is not None
+ ), "Must specify sp_size and sp_group for sequence parallel"
+
+ if "padding_mask" in kwargs:
+ warnings.warn(
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
+ )
+
+ # overwrite attention_mask with padding_mask
+ attention_mask = kwargs.pop("padding_mask")
+ bsz, q_len, _ = hidden_states.size()
+
+ # sp: modify sp_len when sequence parallel mode is ring
+ if sp_mode in ["split_gather", "ring"]:
+ q_len *= sp_size
+
+ query_states = self.q_proj(hidden_states)
+ key_states = self.k_proj(hidden_states)
+ value_states = self.v_proj(hidden_states)
+
+ # sp: all-to-all comminucation when introducing sequence parallel
+ if sp_mode == "all_to_all":
+ query_states = all_to_all_comm(query_states, sp_group)
+ key_states = all_to_all_comm(key_states, sp_group)
+ value_states = all_to_all_comm(value_states, sp_group)
+ bsz, q_len, _ = query_states.size()
+
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+ kv_seq_len = key_states.shape[-2]
+ if past_key_value is not None:
+ if self.layer_idx is None:
+ raise ValueError(
+ f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
+ "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
+ "with a layer index."
+ )
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+
+ # Because the input can be padded, the absolute sequence length depends on the max position id.
+ rotary_seq_len = max(kv_seq_len, position_ids[:, -1].max().item()) + 1
+ cos, sin = self.rotary_emb(value_states, seq_len=rotary_seq_len)
+
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
+
+ use_sliding_windows = (
+ _flash_supports_window_size
+ and getattr(self.config, "sliding_window", None) is not None
+ and kv_seq_len > self.config.sliding_window
+ )
+ if not _flash_supports_window_size:
+ logger.warning_once(
+ "The current flash attention version does not support sliding window attention, for a more memory efficient implementation"
+ " make sure to upgrade flash-attn library."
+ )
+ if past_key_value is not None:
+ # Activate slicing cache only if the config has a value `sliding_windows` attribute
+ cache_has_contents = past_key_value.get_seq_length(self.layer_idx) > 0
+ if (
+ getattr(self.config, "sliding_window", None) is not None
+ and kv_seq_len > self.config.sliding_window
+ and cache_has_contents
+ ):
+ slicing_tokens = 1 - self.config.sliding_window
+
+ past_key = past_key_value[self.layer_idx][0]
+ past_value = past_key_value[self.layer_idx][1]
+
+ past_key = past_key[:, :, slicing_tokens:, :].contiguous()
+ past_value = past_value[:, :, slicing_tokens:, :].contiguous()
+
+ if past_key.shape[-2] != self.config.sliding_window - 1:
+ raise ValueError(
+ f"past key must have a shape of (`batch_size, num_heads, self.config.sliding_window-1, head_dim`), got"
+ f" {past_key.shape}"
+ )
+
+ if attention_mask is not None:
+ attention_mask = attention_mask[:, slicing_tokens:]
+ attention_mask = torch.cat([attention_mask, torch.ones_like(attention_mask[:, -1:])], dim=-1)
+
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+ # repeat k/v heads if n_kv_heads < n_heads
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
+ dropout_rate = 0.0 if not self.training else self.attention_dropout
+
+ # In PEFT, usually we cast the layer norms in float32 for training stability reasons
+ # therefore the input hidden states gets silently casted in float32. Hence, we need
+ # cast them back in float16 just to be sure everything works as expected.
+ input_dtype = query_states.dtype
+ if input_dtype == torch.float32:
+ if torch.is_autocast_enabled():
+ target_dtype = torch.get_autocast_gpu_dtype()
+ # Handle the case where the model is quantized
+ elif hasattr(self.config, "_pre_quantization_dtype"):
+ target_dtype = self.config._pre_quantization_dtype
+ else:
+ target_dtype = self.q_proj.weight.dtype
+
+ logger.warning_once(
+ f"The input hidden states seems to be silently casted in float32, this might be related to"
+ f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
+ f" {target_dtype}."
+ )
+
+ query_states = query_states.to(target_dtype)
+ key_states = key_states.to(target_dtype)
+ value_states = value_states.to(target_dtype)
+ # Reashape to the expected shape for Flash Attention
+ query_states = query_states.transpose(1, 2)
+ key_states = key_states.transpose(1, 2)
+ value_states = value_states.transpose(1, 2)
+ attn_output = self._flash_attention_forward(
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ q_len,
+ dropout=dropout_rate,
+ use_sliding_windows=use_sliding_windows,
+ )
+
+ # sp: all-to-all comminucation when introducing sequence parallel
+ if sp_mode == "all_to_all":
+ attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim).contiguous() # (1, 8, 128)
+ attn_output = all_to_all_comm(attn_output, sp_group, scatter_dim=1, gather_dim=2) # (1, 4, 256)
+ else:
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+
+ attn_output = self.o_proj(attn_output)
+
+ if not output_attentions:
+ attn_weights = None
+ return attn_output, attn_weights, past_key_value
+
+ return forward
+
+
+def get_mixtral_flash_attention_model_forward(shard_config, sp_mode=None, sp_size=None, sp_group=None):
+ logger = logging.get_logger(__name__)
+
+ def forward(
+ self,
+ input_ids: torch.LongTensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ output_router_logits: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, MoeModelOutputWithPast]:
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+ output_router_logits = (
+ output_router_logits if output_router_logits is not None else self.config.output_router_logits
+ )
+ output_hidden_states = (
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+ )
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ # retrieve input_ids and inputs_embeds
+ if input_ids is not None and inputs_embeds is not None:
+ raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
+ elif input_ids is not None:
+ batch_size, seq_length = input_ids.shape
+ elif inputs_embeds is not None:
+ batch_size, seq_length, _ = inputs_embeds.shape
+ else:
+ raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
+
+ past_key_values_length = 0
+
+ if (self.gradient_checkpointing or sp_mode in ["ring", "all_to_all"]) and self.training:
+ if use_cache:
+ logger.warning_once(
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+ )
+ use_cache = False
+ if use_cache:
+ use_legacy_cache = not isinstance(past_key_values, Cache)
+ if use_legacy_cache:
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
+ past_key_values_length = past_key_values.get_usable_length(seq_length)
+
+ if position_ids is None:
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
+ position_ids = torch.arange(
+ past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
+ )
+ position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
+ else:
+ position_ids = position_ids.view(-1, seq_length).long()
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ if attention_mask is not None and self._attn_implementation == "flash_attention_2" and use_cache:
+ is_padding_right = attention_mask[:, -1].sum().item() != batch_size
+ if is_padding_right:
+ raise ValueError(
+ "You are attempting to perform batched generation with padding_side='right'"
+ " this may lead to unexpected behaviour for Flash Attention version of Mixtral. Make sure to "
+ " call `tokenizer.padding_side = 'left'` before tokenizing the input. "
+ )
+ if self._attn_implementation == "flash_attention_2":
+ # 2d mask is passed through the layers
+ attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
+ elif self._attn_implementation == "sdpa" and not output_attentions:
+ # output_attentions=True can not be supported when using SDPA, and we fall back on
+ # the manual implementation that requires a 4D causal mask in all cases.
+ attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
+ attention_mask,
+ (batch_size, seq_length),
+ inputs_embeds,
+ past_key_values_length,
+ )
+ else:
+ # 4d mask is passed through the layers
+ attention_mask = _prepare_4d_causal_attention_mask(
+ attention_mask,
+ (batch_size, seq_length),
+ inputs_embeds,
+ past_key_values_length,
+ sliding_window=self.config.sliding_window,
+ )
+
+ if sp_mode in ["ring", "split_gather"]:
+ inputs_embeds = split_forward_gather_backward(inputs_embeds, 1, sp_group)
+ elif sp_mode == "all_to_all":
+ inputs_embeds = split_forward_gather_backward(inputs_embeds, 1, sp_group, 1 / sp_size)
+ hidden_states = inputs_embeds
+
+ # decoder layers
+ all_hidden_states = () if output_hidden_states else None
+ all_self_attns = () if output_attentions else None
+ all_router_logits = () if output_router_logits else None
+ next_decoder_cache = None
+
+ for decoder_layer in self.layers:
+ if output_hidden_states:
+ all_hidden_states += (hidden_states,)
+
+ if self.gradient_checkpointing and self.training:
+ layer_outputs = self._gradient_checkpointing_func(
+ decoder_layer.__call__,
+ hidden_states,
+ attention_mask,
+ position_ids,
+ past_key_values,
+ output_attentions,
+ output_router_logits,
+ use_cache,
+ )
+ else:
+ layer_outputs = decoder_layer(
+ hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_value=past_key_values,
+ output_attentions=output_attentions,
+ output_router_logits=output_router_logits,
+ use_cache=use_cache,
+ )
+
+ hidden_states = layer_outputs[0]
+
+ if use_cache:
+ next_decoder_cache = layer_outputs[2 if output_attentions else 1]
+
+ if output_attentions:
+ all_self_attns += (layer_outputs[1],)
+
+ if output_router_logits:
+ all_router_logits += (layer_outputs[-1],)
+
+ hidden_states = self.norm(hidden_states)
+
+ if sp_mode == "ring" or sp_mode == "split_gather":
+ hidden_states = gather_forward_split_backward(hidden_states, 1, sp_group)
+ elif sp_mode == "all_to_all":
+ hidden_states = gather_forward_split_backward(hidden_states, 1, sp_group, grad_scale=sp_size)
+
+ # add hidden states from the last decoder layer
+ if output_hidden_states:
+ all_hidden_states += (hidden_states,)
+
+ next_cache = None
+ if use_cache:
+ next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
+
+ if not return_dict:
+ return tuple(
+ v
+ for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_logits]
+ if v is not None
+ )
+ return MoeModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=next_cache,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attns,
+ router_logits=all_router_logits,
+ )
+
+ return forward
diff --git a/colossalai/shardformer/modeling/opt.py b/colossalai/shardformer/modeling/opt.py
index b250b4976..636b46cc4 100644
--- a/colossalai/shardformer/modeling/opt.py
+++ b/colossalai/shardformer/modeling/opt.py
@@ -22,7 +22,7 @@ from colossalai.pipeline.stage_manager import PipelineStageManager
from colossalai.shardformer.layer import ColoAttention
from colossalai.shardformer.shard import ShardConfig
-from ..layer import cross_entropy_1d
+from ..layer import dist_cross_entropy
logger = logging.get_logger(__name__)
@@ -330,30 +330,14 @@ class OPTPipelineForwards:
)
if stage_manager.is_last_stage():
logits = self.lm_head(outputs[0]).contiguous()
- loss = None
- if labels is not None:
- # move labels to correct device to enable model parallelism
- labels = labels.to(logits.device)
- # Shift so that tokens < n predict n
- shift_logits = logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- # Flatten the tokens
-
- if shard_config.enable_tensor_parallelism and shard_config.parallel_output:
- new_vocab_size = logits.shape[-1]
- shift_logits = shift_logits.view(-1, new_vocab_size)
- shift_labels = shift_labels.view(-1)
- loss = cross_entropy_1d(
- shift_logits,
- shift_labels,
- process_group=shard_config.tensor_parallel_process_group,
- vocab_size=self.lm_head.out_features,
- dtype=self.model.decoder.dtype,
- )
- else:
- loss_fct = CrossEntropyLoss()
- shift_logits = shift_logits.view(-1, self.config.vocab_size)
- loss = loss_fct(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1))
+ loss = dist_cross_entropy(
+ labels,
+ logits,
+ shard_config,
+ self.lm_head.out_features,
+ self.config.vocab_size,
+ self.model.decoder.dtype,
+ )
if not return_dict:
output = (logits,) + outputs[1:]
@@ -971,26 +955,9 @@ def get_lm_forward_with_dist_cross_entropy(shard_config: ShardConfig):
)
logits = self.lm_head(outputs[0]).contiguous()
-
- loss = None
- if labels is not None:
- # move labels to correct device to enable model parallelism
- labels = labels.to(logits.device)
- # Shift so that tokens < n predict n
- shift_logits = logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- shift_labels = shift_labels.view(-1)
- # Enable model parallelism
- shift_labels = shift_labels.to(shift_logits.device)
- new_vocab_size = logits.shape[-1]
- shift_logits = shift_logits.view(-1, new_vocab_size)
- loss = cross_entropy_1d(
- shift_logits,
- shift_labels,
- process_group=shard_config.tensor_parallel_process_group,
- vocab_size=self.lm_head.out_features,
- dtype=self.model.decoder.dtype,
- )
+ loss = dist_cross_entropy(
+ labels, logits, shard_config, self.lm_head.out_features, self.config.vocab_size, self.model.decoder.dtype
+ )
if not return_dict:
output = (logits,) + outputs[1:]
diff --git a/colossalai/shardformer/modeling/qwen2.py b/colossalai/shardformer/modeling/qwen2.py
index 11c26822f..538e96c32 100644
--- a/colossalai/shardformer/modeling/qwen2.py
+++ b/colossalai/shardformer/modeling/qwen2.py
@@ -1,6 +1,8 @@
+import math
from typing import List, Optional, Tuple, Union
import torch
+from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
from transformers.modeling_outputs import (
BaseModelOutputWithPast,
@@ -30,9 +32,14 @@ except ImportError:
from transformers.utils import logging
from colossalai.pipeline.stage_manager import PipelineStageManager
+from colossalai.shardformer.layer._operation import (
+ all_to_all_comm,
+ gather_forward_split_backward,
+ split_forward_gather_backward,
+)
from colossalai.shardformer.shard import ShardConfig
-from ..layer import ColoAttention, cross_entropy_1d
+from ..layer import ColoAttention, dist_cross_entropy
class Qwen2PipelineForwards:
@@ -129,7 +136,7 @@ class Qwen2PipelineForwards:
# for the other stages, hidden_states is the output of the previous stage
if shard_config.enable_flash_attention:
# in this case, attention_mask is a dict rather than a tensor
- mask_shape = (batch_size, 1, seq_length_with_past, seq_length_with_past)
+ mask_shape = (batch_size, 1, seq_length, seq_length_with_past)
attention_mask = ColoAttention.prepare_attn_kwargs(
mask_shape,
hidden_states.dtype,
@@ -162,6 +169,21 @@ class Qwen2PipelineForwards:
sliding_window=self.config.sliding_window,
)
+ if shard_config and shard_config.enable_sequence_parallelism:
+ if shard_config.sequence_parallelism_mode in ["split_gather", "ring"]:
+ hidden_states = split_forward_gather_backward(
+ hidden_states,
+ dim=1,
+ process_group=shard_config.tensor_parallel_process_group,
+ )
+ elif shard_config.sequence_parallelism_mode == "all_to_all":
+ hidden_states = split_forward_gather_backward(
+ hidden_states,
+ dim=1,
+ process_group=shard_config.sequence_parallel_process_group,
+ grad_scale=1 / shard_config.sequence_parallel_size,
+ )
+
# decoder layers
all_hidden_states = () if output_hidden_states else None
all_self_attns = () if output_attentions else None
@@ -218,6 +240,20 @@ class Qwen2PipelineForwards:
if stage_manager.is_last_stage():
hidden_states = self.norm(hidden_states)
+ if shard_config and shard_config.enable_sequence_parallelism:
+ if shard_config.sequence_parallelism_mode in ["split_gather", "ring"]:
+ hidden_states = gather_forward_split_backward(
+ hidden_states,
+ dim=1,
+ process_group=shard_config.tensor_parallel_process_group,
+ )
+ elif shard_config.sequence_parallelism_mode == "all_to_all":
+ hidden_states = gather_forward_split_backward(
+ hidden_states,
+ dim=1,
+ process_group=shard_config.sequence_parallel_process_group,
+ grad_scale=shard_config.sequence_parallel_size,
+ )
# add hidden states from the last decoder layer
if output_hidden_states:
all_hidden_states += (hidden_states,)
@@ -317,25 +353,9 @@ class Qwen2PipelineForwards:
if stage_manager.is_last_stage():
hidden_states = outputs[0]
logits = self.lm_head(hidden_states)
- loss = None
- if labels is not None:
- # Shift so that tokens < n predict n
- shift_logits = logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- # Flatten the tokens
- loss_fct = CrossEntropyLoss()
- shift_labels = shift_labels.view(-1)
- # Enable model parallelism
- shift_labels = shift_labels.to(shift_logits.device)
- if shard_config.enable_tensor_parallelism:
- new_vocab_size = logits.shape[-1]
- shift_logits = shift_logits.view(-1, new_vocab_size)
- loss = cross_entropy_1d(
- shift_logits, shift_labels, process_group=shard_config.tensor_parallel_process_group
- )
- else:
- shift_logits = shift_logits.view(-1, self.config.vocab_size)
- loss = loss_fct(shift_logits, shift_labels)
+ loss = dist_cross_entropy(
+ labels, logits, shard_config, self.lm_head.out_features, self.config.vocab_size, logits.dtype
+ )
if not return_dict:
output = (logits,) + outputs[1:]
@@ -469,7 +489,7 @@ class Qwen2PipelineForwards:
return {"hidden_states": hidden_states}
-def get_qwen2_flash_attention_forward(shard_config: ShardConfig):
+def get_qwen2_flash_attention_forward(shard_config: ShardConfig, sp_mode=None, sp_size=None, sp_group=None):
def forward(
self: Qwen2Attention,
hidden_states: torch.Tensor,
@@ -480,11 +500,26 @@ def get_qwen2_flash_attention_forward(shard_config: ShardConfig):
use_cache: bool = False,
**kwargs,
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+ if sp_mode is not None:
+ assert sp_mode in ["all_to_all", "split_gather", "ring"], "Invalid sp_mode"
+ assert (sp_size is not None) and (
+ sp_group is not None
+ ), "Must specify sp_size and sp_group for sequence parallel"
+
bsz, q_len, _ = hidden_states.size()
+ # sp: modify sp_len when sequence parallel mode is ring
+ if sp_mode in ["split_gather", "ring"]:
+ q_len *= sp_size
query_states = self.q_proj(hidden_states)
key_states = self.k_proj(hidden_states)
value_states = self.v_proj(hidden_states)
+ # sp: all-to-all comminucation when introducing sequence parallel
+ if sp_mode == "all_to_all":
+ query_states = all_to_all_comm(query_states, sp_group)
+ key_states = all_to_all_comm(key_states, sp_group)
+ value_states = all_to_all_comm(value_states, sp_group)
+ bsz, q_len, _ = query_states.size()
query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
@@ -538,10 +573,41 @@ def get_qwen2_flash_attention_forward(shard_config: ShardConfig):
key_states = repeat_kv(key_states, self.num_key_value_groups)
value_states = repeat_kv(value_states, self.num_key_value_groups)
- assert isinstance(attention_mask, dict), "Flash Attention Error: attention_mask should be a dict."
- attn_output = ColoAttention.attention(query_states, key_states, value_states, **attention_mask)
+ if shard_config.enable_flash_attention:
+ assert isinstance(attention_mask, dict), "Flash Attention Error: attention_mask should be a dict."
+ attn_output = ColoAttention.attention(query_states, key_states, value_states, **attention_mask)
+ else:
+ attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
+
+ if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
+ raise ValueError(
+ f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
+ f" {attn_weights.size()}"
+ )
+
+ if attention_mask is not None:
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+ raise ValueError(
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
+ )
+ attn_weights = attn_weights + attention_mask
+
+ # upcast attention to fp32
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+ attn_output = torch.matmul(attn_weights, value_states)
+
+ if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
+ raise ValueError(
+ f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
+ f" {attn_output.size()}"
+ )
attn_output = attn_output.transpose(1, 2).contiguous()
- attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+ if sp_mode == "all_to_all":
+ attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim)
+ attn_output = all_to_all_comm(attn_output, sp_group, scatter_dim=1, gather_dim=2)
+ else:
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+
attn_output = self.o_proj(attn_output)
return attn_output, None, past_key_value
@@ -549,9 +615,8 @@ def get_qwen2_flash_attention_forward(shard_config: ShardConfig):
return forward
-def get_qwen2_model_forward_for_flash_attn(shard_config: ShardConfig):
+def get_qwen2_model_forward_for_flash_attn(shard_config: ShardConfig, sp_mode=None, sp_size=None, sp_group=None):
logger = logging.get_logger(__name__)
- assert shard_config.enable_flash_attention, "Flash Attention is not enabled."
def forward(
self,
@@ -586,6 +651,10 @@ def get_qwen2_model_forward_for_flash_attn(shard_config: ShardConfig):
seq_length_with_past = seq_length
past_key_values_length = 0
+ if past_key_values is not None:
+ past_key_values_length = past_key_values[0][0].shape[2]
+ seq_length_with_past = seq_length_with_past + past_key_values_length
+
if position_ids is None:
device = input_ids.device if input_ids is not None else inputs_embeds.device
position_ids = torch.arange(
@@ -601,17 +670,26 @@ def get_qwen2_model_forward_for_flash_attn(shard_config: ShardConfig):
# embed positions
hidden_states = inputs_embeds
- # in this case, attention_mask is a dict rather than a tensor
- mask_shape = (batch_size, 1, seq_length_with_past, seq_length_with_past)
- attention_mask = ColoAttention.prepare_attn_kwargs(
- mask_shape,
- hidden_states.dtype,
- hidden_states.device,
- q_padding_mask=attention_mask,
- is_causal=True,
- )
+ if shard_config.enable_flash_attention:
+ # in this case, attention_mask is a dict rather than a tensor
+ mask_shape = (batch_size, 1, seq_length, seq_length_with_past)
+ attention_mask = ColoAttention.prepare_attn_kwargs(
+ mask_shape,
+ hidden_states.dtype,
+ hidden_states.device,
+ q_padding_mask=attention_mask,
+ is_causal=True,
+ )
+ else:
+ attention_mask = _prepare_4d_causal_attention_mask(
+ attention_mask,
+ (batch_size, seq_length),
+ inputs_embeds,
+ past_key_values_length,
+ sliding_window=self.config.sliding_window,
+ )
- if self.gradient_checkpointing and self.training:
+ if (self.gradient_checkpointing or sp_mode in ["ring", "all_to_all"]) and self.training:
if use_cache:
logger.warning_once(
"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
@@ -623,6 +701,11 @@ def get_qwen2_model_forward_for_flash_attn(shard_config: ShardConfig):
all_self_attns = () if output_attentions else None
next_decoder_cache = None
+ if sp_mode in ["ring", "split_gather"]:
+ hidden_states = split_forward_gather_backward(hidden_states, 1, sp_group)
+ elif sp_mode == "all_to_all":
+ hidden_states = split_forward_gather_backward(hidden_states, 1, sp_group, 1 / sp_size)
+
for decoder_layer in self.layers:
if output_hidden_states:
all_hidden_states += (hidden_states,)
@@ -657,6 +740,11 @@ def get_qwen2_model_forward_for_flash_attn(shard_config: ShardConfig):
hidden_states = self.norm(hidden_states)
+ if sp_mode == "ring" or sp_mode == "split_gather":
+ hidden_states = gather_forward_split_backward(hidden_states, 1, sp_group)
+ elif sp_mode == "all_to_all":
+ hidden_states = gather_forward_split_backward(hidden_states, 1, sp_group, grad_scale=sp_size)
+
# add hidden states from the last decoder layer
if output_hidden_states:
all_hidden_states += (hidden_states,)
@@ -737,26 +825,9 @@ def get_lm_forward_with_dist_cross_entropy(shard_config: ShardConfig):
hidden_states = outputs[0]
logits = self.lm_head(hidden_states)
logits = logits.float()
-
- loss = None
- if labels is not None:
- # Shift so that tokens < n predict n
- shift_logits = logits[..., :-1, :].contiguous()
- shift_labels = labels[..., 1:].contiguous()
- # Flatten the tokens
- loss_fct = CrossEntropyLoss()
- shift_labels = shift_labels.view(-1)
- # Enable model parallelism
- shift_labels = shift_labels.to(shift_logits.device)
- if shard_config.enable_tensor_parallelism:
- new_vocab_size = logits.shape[-1]
- shift_logits = shift_logits.view(-1, new_vocab_size)
- loss = cross_entropy_1d(
- shift_logits, shift_labels, process_group=shard_config.tensor_parallel_process_group
- )
- else:
- shift_logits = shift_logits.view(-1, self.config.vocab_size)
- loss = loss_fct(shift_logits, shift_labels)
+ loss = dist_cross_entropy(
+ labels, logits, shard_config, self.lm_head.out_features, self.config.vocab_size, logits.dtype
+ )
if not return_dict:
output = (logits,) + outputs[1:]
diff --git a/colossalai/shardformer/policies/auto_policy.py b/colossalai/shardformer/policies/auto_policy.py
index bf139c840..7b9c759a6 100644
--- a/colossalai/shardformer/policies/auto_policy.py
+++ b/colossalai/shardformer/policies/auto_policy.py
@@ -160,6 +160,13 @@ _POLICY_LIST = {
"transformers_modules.modeling_chatglm.ChatGLMForConditionalGeneration": PolicyLocation(
file_name="chatglm2", class_name="ChatGLMForConditionalGenerationPolicy"
),
+ # Deepseek
+ "transformers_modules.modeling_deepseek.DeepseekModel": PolicyLocation(
+ file_name="deepseek", class_name="DeepseekModelPolicy"
+ ),
+ "transformers_modules.modeling_deepseek.DeepseekForCausalLM": PolicyLocation(
+ file_name="deepseek", class_name="DeepseekForCausalLMPolicy"
+ ),
# Falcon
"transformers.models.falcon.modeling_falcon.FalconModel": PolicyLocation(
file_name="falcon", class_name="FalconModelPolicy"
@@ -193,6 +200,9 @@ _POLICY_LIST = {
"transformers.models.mixtral.modeling_mixtral.MixtralForCausalLM": PolicyLocation(
file_name="mixtral", class_name="MixtralForCausalLMPolicy"
),
+ "transformers.models.mixtral.modeling_mixtral.MixtralForSequenceClassification": PolicyLocation(
+ file_name="mixtral", class_name="MixtralForSequenceClassificationPolicy"
+ ),
# Qwen2
"transformers.models.qwen2.modeling_qwen2.Qwen2Model": PolicyLocation(
file_name="qwen2", class_name="Qwen2ModelPolicy"
@@ -233,6 +243,9 @@ def _fullname(obj):
# patch custom models which are not in transformers
# it can be like 'transformers_modules.THUDM.chatglm3-6b.103caa40027ebfd8450289ca2f278eac4ff26405.modeling_chatglm' (from huggingface hub)
# or like 'transformers_modules.chatglm.modeling_chatglm' (from local directory)
+ if module.startswith("peft"):
+ klass = obj.base_model.model.__class__
+ module = klass.__module__
if module.startswith("transformers_modules"):
split_module = module.split(".")
if len(split_module) >= 2:
@@ -252,7 +265,6 @@ def get_autopolicy(model: nn.Module) -> Policy:
"""
full_name = _fullname(model)
policy_location = _POLICY_LIST.get(full_name, None)
-
if policy_location is None:
raise NotImplementedError(
f"Auto policy for {model.__class__.__qualname__} ({full_name}) is not implemented\n. Supported models are {list(_POLICY_LIST.keys())}"
diff --git a/colossalai/shardformer/policies/chatglm2.py b/colossalai/shardformer/policies/chatglm2.py
index 01aa77e57..3877bdac3 100644
--- a/colossalai/shardformer/policies/chatglm2.py
+++ b/colossalai/shardformer/policies/chatglm2.py
@@ -9,6 +9,7 @@ import colossalai.shardformer.layer as col_nn
from colossalai.shardformer.modeling.chatglm2 import ChatGLMPipelineForwards
from ..modeling.chatglm2 import (
+ get_chatglm_sequence_parallel_attention_forward,
get_chatglm_sequence_parallel_forward_fn,
get_flash_core_attention_forward,
get_jit_fused_glm_block_forward,
@@ -58,14 +59,29 @@ class ChatGLMPolicy(Policy):
norm_cls = col_nn.LayerNorm
sp_mode = self.shard_config.sequence_parallelism_mode or None
- assert sp_mode != "all_to_all", "all_to_all sequence parallelism is not supported for ChatGLM2"
+ sp_size = self.shard_config.sequence_parallel_size or None
+ sp_group = self.shard_config.sequence_parallel_process_group or None
+
if sp_mode == "ring":
warnings.warn(
f"For ChatGLM2, sequence parallelism is currently not support mode {sp_mode}, will set to be split_gather"
)
sp_mode = "split_gather"
overlap = self.shard_config.enable_sequence_overlap
- sp_partial_derived = sp_mode == "split_gather"
+ sp_partial_derived = sp_mode in ["split_gather"]
+
+ if sp_mode == "all_to_all":
+ decoder_attribute_replacement = {
+ "num_heads": self.model.config.num_attention_heads // sp_size,
+ "hidden_size_per_partition": self.model.config.kv_channels
+ * self.model.config.num_attention_heads
+ // sp_size,
+ }
+ if getattr(self.model.config, "num_key_value_heads", False):
+ decoder_attribute_replacement["num_key_value_heads"] = self.model.config.num_key_value_heads // sp_size
+ policy["CoreAttention"] = ModulePolicyDescription(
+ attribute_replacement=decoder_attribute_replacement,
+ )
if self.shard_config.enable_tensor_parallelism:
assert (
@@ -179,12 +195,26 @@ class ChatGLMPolicy(Policy):
)
# use sequence parallel
- if sp_mode == "split_gather":
+ if self.shard_config.enable_sequence_parallelism:
self.append_or_create_method_replacement(
- description={"forward": get_chatglm_sequence_parallel_forward_fn(self.shard_config)},
+ description={
+ "forward": get_chatglm_sequence_parallel_attention_forward(
+ self.shard_config, sp_mode, sp_size, sp_group
+ ),
+ },
policy=policy,
- target_key="ChatGLMModel",
+ target_key="SelfAttention",
)
+ if self.pipeline_stage_manager is None:
+ self.append_or_create_method_replacement(
+ description={
+ "forward": get_chatglm_sequence_parallel_forward_fn(
+ self.shard_config, sp_mode, sp_size, sp_group
+ )
+ },
+ policy=policy,
+ target_key="ChatGLMModel",
+ )
# use jit fused operator
if self.shard_config.enable_jit_fused:
diff --git a/colossalai/shardformer/policies/command.py b/colossalai/shardformer/policies/command.py
index 902baf2e1..a9b915d10 100644
--- a/colossalai/shardformer/policies/command.py
+++ b/colossalai/shardformer/policies/command.py
@@ -1,4 +1,3 @@
-import warnings
from functools import partial
from typing import Callable, Dict, List, Union
@@ -66,13 +65,6 @@ class CommandPolicy(Policy):
else:
norm_cls = LayerNorm
- if self.pipeline_stage_manager is not None:
- self.shard_config.enable_sequence_parallelism = False
- self.shard_config.enable_sequence_overlap = False
- self.shard_config.sequence_parallelism_mode = None
- warnings.warn(
- f"For Command, sequence parallelism is currently not compatible with pipeline parallelism, set to be False"
- )
sp_mode = self.shard_config.sequence_parallelism_mode or None
sp_size = self.shard_config.sequence_parallel_size or None
sp_group = self.shard_config.sequence_parallel_process_group or None
diff --git a/colossalai/shardformer/policies/deepseek.py b/colossalai/shardformer/policies/deepseek.py
new file mode 100644
index 000000000..605f69c4a
--- /dev/null
+++ b/colossalai/shardformer/policies/deepseek.py
@@ -0,0 +1,347 @@
+from functools import partial
+from typing import Callable, Dict, List, Union
+
+import torch.nn as nn
+from torch import Tensor
+from torch.nn import Module
+from transformers.utils import is_flash_attn_greater_or_equal_2_10
+
+from colossalai.shardformer.layer import FusedRMSNorm, Linear1D_Col
+from colossalai.shardformer.layer.embedding import PaddingEmbedding, VocabParallelEmbedding1D
+from colossalai.shardformer.layer.linear import Linear1D_Row
+from colossalai.shardformer.modeling.deepseek import (
+ DeepseekPipelineForwards,
+ EPDeepseekMoE,
+ get_deepseek_flash_attention_forward,
+ get_deepseek_flash_attention_model_forward,
+)
+from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription
+
+__all__ = ["DeepseekPolicy", "DeepseekForCausalLMPolicy"]
+
+
+class DeepseekPolicy(Policy):
+ def config_sanity_check(self):
+ pass
+
+ def preprocess(self):
+ self.tie_weight = self.tie_weight_check()
+ self.origin_attn_implement = self.model.config._attn_implementation
+ """
+ Because transformers library's bug for AutoModel/AutoConfig, who pop “attn_implement” twice from modeling_utils.py and configuration_utils.py.
+ This bug causes attn_cls to be set to sdpa. Here we assign it to "flash_attention_2".
+ """
+ # self.origin_attn_implement = "flash_attention_2"
+ if self.shard_config.enable_tensor_parallelism:
+ # Resize embedding
+ vocab_size = self.model.config.vocab_size
+ world_size = self.shard_config.tensor_parallel_size
+
+ if vocab_size % world_size != 0:
+ new_vocab_size = vocab_size + world_size - vocab_size % world_size
+ self.model.resize_token_embeddings(new_vocab_size)
+
+ return self.model
+
+ def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]:
+
+ ATTN_IMPLEMENTATION = {
+ "eager": "DeepseekAttention",
+ "flash_attention_2": "DeepseekFlashAttention2",
+ "sdpa": "DeepseekSdpaAttention",
+ }
+ policy = {}
+ attn_cls = ATTN_IMPLEMENTATION[self.origin_attn_implement]
+ sp_mode = self.shard_config.sequence_parallelism_mode or None
+ sp_size = self.shard_config.sequence_parallel_size or None
+ sp_group = self.shard_config.sequence_parallel_process_group or None
+ sp_partial_derived = sp_mode in ["split_gather", "ring"]
+ if sp_mode == "all_to_all":
+ decoder_attribute_replacement = {
+ "num_heads": self.model.config.num_attention_heads // sp_size,
+ }
+ if getattr(self.model.config, "num_key_value_heads", False):
+ decoder_attribute_replacement["num_key_value_heads"] = self.model.config.num_key_value_heads // sp_size
+
+ policy[attn_cls] = ModulePolicyDescription(
+ attribute_replacement=decoder_attribute_replacement,
+ )
+ if self.shard_config.enable_sequence_parallelism:
+ if self.pipeline_stage_manager is not None:
+ # NOTE: we are replacing model forward for both sequence parallelism and pipeline parallelism
+ # if both are enabled, one of them will be ignored
+ raise NotImplementedError("Sequence parallelism is not supported with pipeline parallelism.")
+ self.append_or_create_method_replacement(
+ description={
+ "forward": get_deepseek_flash_attention_forward(self.shard_config, sp_mode, sp_size, sp_group),
+ },
+ policy=policy,
+ target_key=attn_cls,
+ )
+ if self.pipeline_stage_manager is None:
+ self.append_or_create_method_replacement(
+ description={
+ "forward": get_deepseek_flash_attention_model_forward(
+ self.shard_config,
+ sp_mode=sp_mode,
+ sp_size=sp_size,
+ sp_group=sp_group,
+ ),
+ },
+ policy=policy,
+ target_key="DeepseekModel",
+ )
+ embedding_cls = None
+ if self.shard_config.enable_tensor_parallelism:
+ embedding_cls = VocabParallelEmbedding1D
+ else:
+ if self.tie_weight:
+ embedding_cls = PaddingEmbedding
+ if self.shard_config.enable_tensor_parallelism:
+ # tensor parallelism for non-moe params
+ assert (
+ self.model.config.num_attention_heads % self.shard_config.tensor_parallel_size == 0
+ ), f"The number of attention heads must be divisible by tensor parallel size."
+ assert (
+ self.model.config.num_key_value_heads % self.shard_config.tensor_parallel_size == 0
+ ), f"The number of key_value heads must be divisible by tensor parallel size."
+ decoder_attribute_replacement = {
+ "self_attn.hidden_size": self.model.config.hidden_size // self.shard_config.tensor_parallel_size,
+ "self_attn.num_heads": self.model.config.num_attention_heads // self.shard_config.tensor_parallel_size,
+ "self_attn.num_key_value_heads": self.model.config.num_key_value_heads
+ // self.shard_config.tensor_parallel_size,
+ }
+
+ policy["DeepseekDecoderLayer"] = ModulePolicyDescription(
+ attribute_replacement=decoder_attribute_replacement,
+ sub_module_replacement=[
+ SubModuleReplacementDescription(
+ suffix="self_attn.q_proj",
+ target_module=Linear1D_Col,
+ ),
+ SubModuleReplacementDescription(
+ suffix="self_attn.k_proj",
+ target_module=Linear1D_Col,
+ ),
+ SubModuleReplacementDescription(
+ suffix="self_attn.v_proj",
+ target_module=Linear1D_Col,
+ ),
+ SubModuleReplacementDescription(
+ suffix="self_attn.o_proj",
+ target_module=Linear1D_Row,
+ ),
+ ],
+ )
+ if embedding_cls is not None:
+ self.append_or_create_submodule_replacement(
+ description=SubModuleReplacementDescription(
+ suffix="embed_tokens",
+ target_module=embedding_cls,
+ kwargs={"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by},
+ ),
+ policy=policy,
+ target_key="DeepseekModel",
+ )
+
+ if self.shard_config.ep_group:
+ # expert parallel
+ self.append_or_create_submodule_replacement(
+ description=[
+ SubModuleReplacementDescription(
+ suffix="mlp",
+ target_module=EPDeepseekMoE,
+ kwargs={
+ "ep_group": self.shard_config.ep_group,
+ "tp_group": self.shard_config.tensor_parallel_process_group,
+ "moe_dp_group": self.shard_config.moe_dp_group,
+ },
+ )
+ ],
+ policy=policy,
+ target_key="DeepseekDecoderLayer",
+ )
+
+ # optimization configuration
+ if self.shard_config.enable_fused_normalization:
+ self.append_or_create_submodule_replacement(
+ description=[
+ SubModuleReplacementDescription(
+ suffix="input_layernorm",
+ target_module=FusedRMSNorm,
+ kwargs={"sp_partial_derived": sp_partial_derived},
+ ),
+ SubModuleReplacementDescription(
+ suffix="post_attention_layernorm",
+ target_module=FusedRMSNorm,
+ kwargs={"sp_partial_derived": sp_partial_derived},
+ ),
+ ],
+ policy=policy,
+ target_key="DeepseekDecoderLayer",
+ )
+
+ self.append_or_create_submodule_replacement(
+ description=SubModuleReplacementDescription(
+ suffix="norm",
+ target_module=FusedRMSNorm,
+ kwargs={"sp_partial_derived": sp_partial_derived},
+ ),
+ policy=policy,
+ target_key="DeepseekModel",
+ )
+
+ if self.shard_config.enable_flash_attention:
+ # NOTE: there is a bug for toggling flash attention in AutoModel, which has to be used for deepseek right now
+ from transformers.dynamic_module_utils import get_class_from_dynamic_module
+
+ flash_attn_cls = get_class_from_dynamic_module(
+ "deepseek-ai/deepseek-moe-16b-base--modeling_deepseek.DeepseekFlashAttention2",
+ "deepseek-ai/deepseek-moe-16b-base",
+ )
+
+ class TargetFlashAttn:
+ def __init__(self):
+ raise RuntimeError("This class should not be instantiated")
+
+ @staticmethod
+ def from_native_module(original_attn: nn.Module, *args, **kwargs) -> nn.Module:
+ original_attn.__class__ = flash_attn_cls
+ original_attn._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
+ return original_attn
+
+ self.append_or_create_submodule_replacement(
+ description=SubModuleReplacementDescription(
+ suffix="self_attn",
+ target_module=TargetFlashAttn,
+ ),
+ policy=policy,
+ target_key="DeepseekDecoderLayer",
+ )
+ return policy
+
+ def postprocess(self):
+ return self.model
+
+ def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, policy: Dict) -> None:
+ """If under pipeline parallel setting, replacing the original forward method of huggingface
+ to customized forward method, and add this changing to policy."""
+ if self.pipeline_stage_manager:
+ if self.shard_config.enable_sequence_parallelism:
+ # NOTE: we are replacing model forward for both sequence parallelism and pipeline parallelism
+ # if both are enabled, one of them will be ignored
+ raise NotImplementedError("Pipeline parallelism is not supported with sequence parallelism.")
+ stage_manager = self.pipeline_stage_manager
+ if self.model.__class__.__name__ == "DeepseekModel":
+ module = self.model
+ else:
+ module = self.model.model
+
+ layers_per_stage = stage_manager.distribute_layers(len(module.layers))
+ stage_index = stage_manager.get_stage_index(layers_per_stage)
+ method_replacement = {"forward": partial(new_forward, stage_manager=stage_manager, stage_index=stage_index)}
+ self.append_or_create_method_replacement(
+ description=method_replacement, policy=policy, target_key=model_cls
+ )
+
+ return
+
+ def get_held_layers(self) -> List[Module]:
+ """Get pipeline layers for current stage."""
+ assert self.pipeline_stage_manager is not None
+
+ if self.model.__class__.__name__ == "DeepseekModel":
+ module = self.model
+ else:
+ module = self.model.model
+ stage_manager = self.pipeline_stage_manager
+
+ held_layers = []
+ layers_per_stage = stage_manager.distribute_layers(len(module.layers))
+ if stage_manager.is_first_stage():
+ held_layers.append(module.embed_tokens)
+ start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage)
+ held_layers.extend(module.layers[start_idx:end_idx])
+ if stage_manager.is_last_stage():
+ held_layers.append(module.norm)
+
+ return held_layers
+
+
+class DeepseekModelPolicy(DeepseekPolicy):
+ def __init__(self) -> None:
+ super().__init__()
+
+ def module_policy(self):
+ policy = super().module_policy()
+ if self.pipeline_stage_manager:
+ # set None as default
+ self.set_pipeline_forward(
+ model_cls="DeepseekModel",
+ new_forward=DeepseekPipelineForwards.deepseek_model_forward,
+ policy=policy,
+ )
+ return policy
+
+ def get_held_layers(self) -> List[Module]:
+ """Get pipeline layers for current stage."""
+ held_layers = super().get_held_layers()
+ return held_layers
+
+ def get_shared_params(self) -> List[Dict[int, Tensor]]:
+ """No shared params in llama model"""
+ return []
+
+
+class DeepseekForCausalLMPolicy(DeepseekPolicy):
+ def module_policy(self):
+ policy = super().module_policy()
+ # TODO: assign pg mesh from plugin to all modules
+ if self.shard_config.enable_tensor_parallelism:
+ # add a new item for casual lm
+ new_item = {
+ "DeepseekForCausalLM": ModulePolicyDescription(
+ sub_module_replacement=[
+ SubModuleReplacementDescription(
+ suffix="lm_head",
+ target_module=Linear1D_Col,
+ kwargs=dict(gather_output=True),
+ )
+ ]
+ )
+ }
+ policy.update(new_item)
+
+ if self.pipeline_stage_manager:
+ # set None as default
+ self.set_pipeline_forward(
+ model_cls="DeepseekForCausalLM",
+ new_forward=DeepseekPipelineForwards.deepseek_for_causal_lm_forward,
+ policy=policy,
+ )
+
+ return policy
+
+ def get_held_layers(self) -> List[Module]:
+ """Get pipeline layers for current stage."""
+ stage_manager = self.pipeline_stage_manager
+ held_layers = super().get_held_layers()
+ if stage_manager.is_last_stage():
+ held_layers.append(self.model.lm_head)
+ return held_layers
+
+ def get_shared_params(self) -> List[Dict[int, Tensor]]:
+ deepseek_model = self.model.model
+ if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1:
+ if (
+ id(deepseek_model.embed_tokens.weight) == id(self.model.lm_head.weight)
+ and self.pipeline_stage_manager.num_stages > 1
+ ):
+ # tie weights
+ return [
+ {
+ 0: deepseek_model.embed_tokens.weight,
+ self.pipeline_stage_manager.num_stages - 1: self.model.lm_head.weight,
+ }
+ ]
+ return []
diff --git a/colossalai/shardformer/policies/llama.py b/colossalai/shardformer/policies/llama.py
index 0f28f6cf4..6f8404219 100644
--- a/colossalai/shardformer/policies/llama.py
+++ b/colossalai/shardformer/policies/llama.py
@@ -1,4 +1,3 @@
-import warnings
from functools import partial
from typing import Callable, Dict, List, Union
@@ -65,13 +64,7 @@ class LlamaPolicy(Policy):
norm_cls = FusedRMSNorm
else:
norm_cls = RMSNorm
- if self.pipeline_stage_manager is not None:
- self.shard_config.enable_sequence_parallelism = False
- self.shard_config.enable_sequence_overlap = False
- self.shard_config.sequence_parallelism_mode = None
- warnings.warn(
- f"For llama, sequence parallelism is currently not compatible with pipeline parallelism, set to be False"
- )
+
sp_mode = self.shard_config.sequence_parallelism_mode or None
sp_size = self.shard_config.sequence_parallel_size or None
sp_group = self.shard_config.sequence_parallel_process_group or None
diff --git a/colossalai/shardformer/policies/mixtral.py b/colossalai/shardformer/policies/mixtral.py
index 0fb858d78..10df143c9 100644
--- a/colossalai/shardformer/policies/mixtral.py
+++ b/colossalai/shardformer/policies/mixtral.py
@@ -1,13 +1,21 @@
+import warnings
from functools import partial
from typing import Callable, Dict, List, Union
import torch.nn as nn
from torch import Tensor
from torch.nn import Module
-from transformers.models.mixtral.modeling_mixtral import MixtralDecoderLayer, MixtralForCausalLM, MixtralModel
+from transformers.models.mixtral.modeling_mixtral import MixtralForCausalLM, MixtralModel
from colossalai.shardformer.layer import FusedRMSNorm, Linear1D_Col
-from colossalai.shardformer.modeling.mixtral import EPMixtralSparseMoeBlock, MixtralPipelineForwards
+from colossalai.shardformer.layer.embedding import PaddingEmbedding, VocabParallelEmbedding1D
+from colossalai.shardformer.layer.linear import Linear1D_Row
+from colossalai.shardformer.modeling.mixtral import (
+ EPMixtralSparseMoeBlock,
+ MixtralPipelineForwards,
+ get_mixtral_flash_attention_forward,
+ get_mixtral_flash_attention_model_forward,
+)
from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription
__all__ = ["MixtralPolicy", "MixtralForCausalLMPolicy"]
@@ -18,36 +26,136 @@ class MixtralPolicy(Policy):
pass
def preprocess(self):
- if self.shard_config.enable_tensor_parallelism:
- # Resize embedding
- vocab_size = self.model.config.vocab_size
- world_size = self.shard_config.tensor_parallel_size
-
- if vocab_size % world_size != 0:
- new_vocab_size = vocab_size + world_size - vocab_size % world_size
- self.model.resize_token_embeddings(new_vocab_size)
-
+ self.tie_weight = self.tie_weight_check()
+ self.origin_attn_implement = self.model.config._attn_implementation
return self.model
def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]:
- policy = {}
+ from transformers.models.mixtral.modeling_mixtral import (
+ MixtralAttention,
+ MixtralDecoderLayer,
+ MixtralFlashAttention2,
+ MixtralModel,
+ MixtralSdpaAttention,
+ )
+ ATTN_IMPLEMENTATION = {
+ "eager": MixtralAttention,
+ "flash_attention_2": MixtralFlashAttention2,
+ "sdpa": MixtralSdpaAttention,
+ }
+ policy = {}
+ attn_cls = ATTN_IMPLEMENTATION[self.origin_attn_implement]
+
+ sp_mode = self.shard_config.sequence_parallelism_mode or None
+ sp_size = self.shard_config.sequence_parallel_size or None
+ sp_group = self.shard_config.sequence_parallel_process_group or None
+ sp_partial_derived = sp_mode in ["split_gather", "ring"]
+ if sp_mode == "all_to_all":
+ decoder_attribute_replacement = {
+ "num_heads": self.model.config.num_attention_heads // sp_size,
+ }
+ if getattr(self.model.config, "num_key_value_heads", False):
+ decoder_attribute_replacement["num_key_value_heads"] = self.model.config.num_key_value_heads // sp_size
+
+ policy[attn_cls] = ModulePolicyDescription(
+ attribute_replacement=decoder_attribute_replacement,
+ )
if self.shard_config.enable_sequence_parallelism:
- self.shard_config.enable_sequence_parallelism = False
- raise NotImplementedError(
- "Mixtral dosen't support sequence parallelism now, will ignore the sequence parallelism flag."
+ if self.pipeline_stage_manager is not None:
+ # NOTE: we are replacing model forward for both sequence parallelism and pipeline parallelism
+ # if both are enabled, one of them will be ignored
+ raise NotImplementedError("Sequence parallelism is not supported with pipeline parallelism.")
+ self.append_or_create_method_replacement(
+ description={
+ "forward": get_mixtral_flash_attention_forward(self.shard_config, sp_mode, sp_size, sp_group),
+ },
+ policy=policy,
+ target_key=attn_cls,
+ )
+ self.append_or_create_method_replacement(
+ description={
+ "forward": get_mixtral_flash_attention_model_forward(
+ self.shard_config,
+ sp_mode=sp_mode,
+ sp_size=sp_size,
+ sp_group=sp_group,
+ ),
+ },
+ policy=policy,
+ target_key=MixtralModel,
)
+ embedding_cls = None
if self.shard_config.enable_tensor_parallelism:
- raise NotImplementedError("Tensor parallelism is not supported for Mixtral model now.")
- if getattr(self.shard_config, "ep_group", None) is not None:
+ embedding_cls = VocabParallelEmbedding1D
+ else:
+ if self.tie_weight:
+ embedding_cls = PaddingEmbedding
+
+ if self.shard_config.enable_tensor_parallelism:
+ # tensor parallelism for non-moe params
+ assert (
+ self.model.config.num_attention_heads % self.shard_config.tensor_parallel_size == 0
+ ), f"The number of attention heads must be divisible by tensor parallel size."
+ assert (
+ self.model.config.num_key_value_heads % self.shard_config.tensor_parallel_size == 0
+ ), f"The number of key_value heads must be divisible by tensor parallel size."
+ decoder_attribute_replacement = {
+ "self_attn.hidden_size": self.model.config.hidden_size // self.shard_config.tensor_parallel_size,
+ "self_attn.num_heads": self.model.config.num_attention_heads // self.shard_config.tensor_parallel_size,
+ "self_attn.num_key_value_heads": self.model.config.num_key_value_heads
+ // self.shard_config.tensor_parallel_size,
+ }
+
+ policy[MixtralDecoderLayer] = ModulePolicyDescription(
+ attribute_replacement=decoder_attribute_replacement,
+ sub_module_replacement=[
+ SubModuleReplacementDescription(
+ suffix="self_attn.q_proj",
+ target_module=Linear1D_Col,
+ ),
+ SubModuleReplacementDescription(
+ suffix="self_attn.k_proj",
+ target_module=Linear1D_Col,
+ ),
+ SubModuleReplacementDescription(
+ suffix="self_attn.v_proj",
+ target_module=Linear1D_Col,
+ ),
+ SubModuleReplacementDescription(
+ suffix="self_attn.o_proj",
+ target_module=Linear1D_Row,
+ ),
+ SubModuleReplacementDescription( # or replicate?
+ suffix="block_sparse_moe.gate", target_module=Linear1D_Col, kwargs={"gather_output": True}
+ ),
+ ],
+ )
+
+ if embedding_cls is not None:
+ self.append_or_create_submodule_replacement(
+ description=SubModuleReplacementDescription(
+ suffix="embed_tokens",
+ target_module=embedding_cls,
+ kwargs={"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by},
+ ),
+ policy=policy,
+ target_key=MixtralModel,
+ )
+
+ if self.shard_config.ep_group:
# expert parallel
self.append_or_create_submodule_replacement(
description=[
SubModuleReplacementDescription(
suffix="block_sparse_moe",
target_module=EPMixtralSparseMoeBlock,
- kwargs={"ep_group": self.shard_config.ep_group},
+ kwargs={
+ "ep_group": self.shard_config.ep_group,
+ "tp_group": self.shard_config.tensor_parallel_process_group,
+ "moe_dp_group": self.shard_config.moe_dp_group,
+ },
)
],
policy=policy,
@@ -61,10 +169,12 @@ class MixtralPolicy(Policy):
SubModuleReplacementDescription(
suffix="input_layernorm",
target_module=FusedRMSNorm,
+ kwargs={"sp_partial_derived": sp_partial_derived},
),
SubModuleReplacementDescription(
suffix="post_attention_layernorm",
target_module=FusedRMSNorm,
+ kwargs={"sp_partial_derived": sp_partial_derived},
),
],
policy=policy,
@@ -75,13 +185,15 @@ class MixtralPolicy(Policy):
description=SubModuleReplacementDescription(
suffix="norm",
target_module=FusedRMSNorm,
+ kwargs={"sp_partial_derived": sp_partial_derived},
),
policy=policy,
target_key=MixtralModel,
)
if self.shard_config.enable_flash_attention:
- raise NotImplementedError("Flash attention has already been replaced in mixtral.")
+ warnings.warn("Flash attention is natively supported in transformers, will ignore the flag.")
+ self.shard_config.enable_flash_attention = False
return policy
@@ -92,6 +204,10 @@ class MixtralPolicy(Policy):
"""If under pipeline parallel setting, replacing the original forward method of huggingface
to customized forward method, and add this changing to policy."""
if self.pipeline_stage_manager:
+ if self.shard_config.enable_sequence_parallelism:
+ # NOTE: we are replacing model forward for both sequence parallelism and pipeline parallelism
+ # if both are enabled, one of them will be ignored
+ raise NotImplementedError("Pipeline parallelism is not supported with sequence parallelism.")
stage_manager = self.pipeline_stage_manager
if self.model.__class__.__name__ == "MixtralModel":
module = self.model
@@ -150,7 +266,7 @@ class MixtralModelPolicy(MixtralPolicy):
return held_layers
def get_shared_params(self) -> List[Dict[int, Tensor]]:
- """No shared params in llama model"""
+ """No shared params in mixtral model"""
return []
@@ -192,17 +308,54 @@ class MixtralForCausalLMPolicy(MixtralPolicy):
return held_layers
def get_shared_params(self) -> List[Dict[int, Tensor]]:
- llama_model = self.model.model
+ mixtral_model = self.model.model
if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1:
if (
- id(llama_model.embed_tokens.weight) == id(self.model.lm_head.weight)
+ id(mixtral_model.embed_tokens.weight) == id(self.model.lm_head.weight)
and self.pipeline_stage_manager.num_stages > 1
):
# tie weights
return [
{
- 0: llama_model.embed_tokens.weight,
+ 0: mixtral_model.embed_tokens.weight,
self.pipeline_stage_manager.num_stages - 1: self.model.lm_head.weight,
}
]
return []
+
+
+class MixtralForSequenceClassificationPolicy(MixtralPolicy):
+ def module_policy(self):
+ from transformers import MixtralForSequenceClassification
+
+ policy = super().module_policy()
+
+ if self.shard_config.enable_tensor_parallelism:
+ # add a new item for sequence classification
+ new_item = {
+ MixtralForSequenceClassification: ModulePolicyDescription(
+ sub_module_replacement=[
+ SubModuleReplacementDescription(
+ suffix="score", target_module=Linear1D_Col, kwargs=dict(gather_output=True)
+ )
+ ]
+ )
+ }
+ policy.update(new_item)
+
+ if self.pipeline_stage_manager:
+ raise NotImplementedError
+
+ return policy
+
+ def get_held_layers(self) -> List[Module]:
+ """Get pipeline layers for current stage."""
+ stage_manager = self.pipeline_stage_manager
+ held_layers = super().get_held_layers()
+ if stage_manager.is_last_stage(ignore_chunk=True):
+ held_layers.append(self.model.score)
+ return held_layers
+
+ def get_shared_params(self) -> List[Dict[int, Tensor]]:
+ """No shared params in mixtral for sequence classification model"""
+ return []
diff --git a/colossalai/shardformer/policies/qwen2.py b/colossalai/shardformer/policies/qwen2.py
index 3e427c4a1..362c14060 100644
--- a/colossalai/shardformer/policies/qwen2.py
+++ b/colossalai/shardformer/policies/qwen2.py
@@ -1,4 +1,3 @@
-import warnings
from functools import partial
from typing import Callable, Dict, List, Union
@@ -82,9 +81,20 @@ class Qwen2Policy(Policy):
embedding_cls = PaddingEmbedding
norm_cls = FusedRMSNorm if self.shard_config.enable_fused_normalization else RMSNorm
- if self.shard_config.enable_sequence_parallelism:
- self.shard_config.enable_sequence_parallelism = False
- warnings.warn("Qwen2 doesn't support sequence parallelism now, will ignore the sequence parallelism flag.")
+ sp_mode = self.shard_config.sequence_parallelism_mode or None
+ sp_size = self.shard_config.sequence_parallel_size or None
+ sp_group = self.shard_config.sequence_parallel_process_group or None
+ sp_partial_derived = sp_mode in ["split_gather", "ring"]
+ if sp_mode == "all_to_all":
+ decoder_attribute_replacement = {
+ "num_heads": self.model.config.num_attention_heads // sp_size,
+ }
+ if getattr(self.model.config, "num_key_value_heads", False):
+ decoder_attribute_replacement["num_key_value_heads"] = self.model.config.num_key_value_heads // sp_size
+
+ policy[attn_cls] = ModulePolicyDescription(
+ attribute_replacement=decoder_attribute_replacement,
+ )
if self.shard_config.enable_tensor_parallelism:
assert (
@@ -109,30 +119,37 @@ class Qwen2Policy(Policy):
SubModuleReplacementDescription(
suffix="self_attn.q_proj",
target_module=Linear1D_Col,
+ kwargs=dict(seq_parallel_mode=sp_mode),
),
SubModuleReplacementDescription(
suffix="self_attn.k_proj",
target_module=Linear1D_Col,
+ kwargs=dict(seq_parallel_mode=sp_mode),
),
SubModuleReplacementDescription(
suffix="self_attn.v_proj",
target_module=Linear1D_Col,
+ kwargs=dict(seq_parallel_mode=sp_mode),
),
SubModuleReplacementDescription(
suffix="self_attn.o_proj",
target_module=Linear1D_Row,
+ kwargs=dict(seq_parallel_mode=sp_mode),
),
SubModuleReplacementDescription(
suffix="mlp.gate_proj",
target_module=Linear1D_Col,
+ kwargs=dict(seq_parallel_mode=sp_mode),
),
SubModuleReplacementDescription(
suffix="mlp.up_proj",
target_module=Linear1D_Col,
+ kwargs=dict(seq_parallel_mode=sp_mode),
),
SubModuleReplacementDescription(
suffix="mlp.down_proj",
target_module=Linear1D_Row,
+ kwargs=dict(seq_parallel_mode=sp_mode),
),
],
)
@@ -154,10 +171,12 @@ class Qwen2Policy(Policy):
SubModuleReplacementDescription(
suffix="input_layernorm",
target_module=norm_cls,
+ kwargs={"sp_partial_derived": sp_partial_derived},
),
SubModuleReplacementDescription(
suffix="post_attention_layernorm",
target_module=norm_cls,
+ kwargs={"sp_partial_derived": sp_partial_derived},
),
],
policy=policy,
@@ -168,16 +187,16 @@ class Qwen2Policy(Policy):
description=SubModuleReplacementDescription(
suffix="norm",
target_module=norm_cls,
+ kwargs={"sp_partial_derived": sp_partial_derived},
),
policy=policy,
target_key=Qwen2Model,
)
- # use flash attention
- if self.shard_config.enable_flash_attention:
+ if self.shard_config.enable_flash_attention or self.shard_config.enable_sequence_parallelism:
self.append_or_create_method_replacement(
description={
- "forward": get_qwen2_flash_attention_forward(self.shard_config),
+ "forward": get_qwen2_flash_attention_forward(self.shard_config, sp_mode, sp_size, sp_group),
},
policy=policy,
target_key=attn_cls,
@@ -186,7 +205,9 @@ class Qwen2Policy(Policy):
# replace qwen2 model forward method
self.append_or_create_method_replacement(
description={
- "forward": get_qwen2_model_forward_for_flash_attn(self.shard_config),
+ "forward": get_qwen2_model_forward_for_flash_attn(
+ self.shard_config, sp_mode, sp_size, sp_group
+ ),
},
policy=policy,
target_key=Qwen2Model,
diff --git a/colossalai/shardformer/shard/shard_config.py b/colossalai/shardformer/shard/shard_config.py
index 7372e06c2..8e33d786b 100644
--- a/colossalai/shardformer/shard/shard_config.py
+++ b/colossalai/shardformer/shard/shard_config.py
@@ -47,6 +47,9 @@ class ShardConfig:
make_vocab_size_divisible_by: int = 64
gradient_checkpoint_config: Optional[GradientCheckpointConfig] = None
extra_kwargs: Dict[str, Any] = field(default_factory=dict)
+
+ # for moe related
+ moe_dp_group: Optional[ProcessGroup] = None
ep_group: Optional[ProcessGroup] = None
fp8_communication: bool = False
# pipeline_parallel_size: int
diff --git a/colossalai/shardformer/shard/shardformer.py b/colossalai/shardformer/shard/shardformer.py
index b54c58273..db03eec41 100644
--- a/colossalai/shardformer/shard/shardformer.py
+++ b/colossalai/shardformer/shard/shardformer.py
@@ -1,4 +1,3 @@
-import os
from typing import Dict, List, Tuple
import torch.distributed as dist
@@ -11,9 +10,6 @@ from ..policies.base_policy import Policy
from .shard_config import ShardConfig
from .sharder import ModelSharder
-# set CUDA_DEVICE_MAX_CONNECTIONS=1 to ensure that when communication and computation overlap, the order of core scheduling is correct
-os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "1"
-
class ShardFormer:
"""
diff --git a/colossalai/tensor/d_tensor/layout_converter.py b/colossalai/tensor/d_tensor/layout_converter.py
index c2cf73181..0f0150d90 100644
--- a/colossalai/tensor/d_tensor/layout_converter.py
+++ b/colossalai/tensor/d_tensor/layout_converter.py
@@ -473,7 +473,7 @@ class LayoutConverter(metaclass=SingletonMeta):
for process_group in used_process_groups:
try:
dist.get_rank(process_group)
- except RuntimeError as e:
+ except (ValueError, RuntimeError) as e:
# If the group is not registered, it means it has been deleted
if str(e) == (
f"Group {process_group} is not registered, please create group with torch.distributed.new_group API"
diff --git a/colossalai/tensor/d_tensor/sharding_spec.py b/colossalai/tensor/d_tensor/sharding_spec.py
index 16a4f248b..76d85a112 100644
--- a/colossalai/tensor/d_tensor/sharding_spec.py
+++ b/colossalai/tensor/d_tensor/sharding_spec.py
@@ -1,4 +1,3 @@
-from copy import deepcopy
from typing import Dict, List
from ..utils import merge_same_dim_mesh_list
@@ -23,10 +22,11 @@ class DimSpec:
Otherwise, the element in shard_list means the data will be sharded in that dimension.
"""
+ _DIFFERENCE_DICT = None
+
def __init__(self, shard_list):
self.is_replica = len(shard_list) == 0
self.shard_list = shard_list
- self.build_difference_2d_dict()
def __eq__(self, other):
return str(self) == str(other)
@@ -39,24 +39,43 @@ class DimSpec:
target += str(dim)
return target
- def _convert_str_to_shard_list(self, str_spec):
+ @property
+ def difference_dict(self):
"""
- Convert str_spec into shard_list.
+ Returns the difference dict, and lazily initializes it when needed
+
+ Return:
+ difference_dict(Dict[Tuple[int, int], Union[int, float, str]]):
+ difference dict
+ """
+ if self._DIFFERENCE_DICT is None:
+ self._DIFFERENCE_DICT = self._build_difference_2d_dict()
+
+ return self._DIFFERENCE_DICT
+
+ def dim_diff(self, other):
+ """
+ The difference between two DimSpec.
Argument:
- str_spec(str): dim spec in str type.
+ other(DimSpec): the dim spec to compare with.
+
+ Return:
+ difference(int): the difference between two DimSpec.
+
+ Example:
+ dim_spec = DimSpec([0])
+ other_dim_spec = DimSpec([0, 1])
+ print(dim_spec.dim_diff(other_dim_spec))
+
+ Output:
+ 5
"""
+ difference = self.difference_dict[(str(self), str(other))]
+ return difference
- if str_spec == "R":
- return []
- if str_spec == "S0":
- return [0]
- if str_spec == "S1":
- return [1]
- if str_spec == "S01":
- return [0, 1]
-
- def build_difference_2d_dict(self):
+ @classmethod
+ def _build_difference_2d_dict(cls):
"""
Build a difference mapping for 2D device mesh case. It will be used to
compute the difference between DimSpec pairs.
@@ -67,9 +86,8 @@ class DimSpec:
difference_dict = {}
for source_spec in source_spec_list:
for target_spec in target_spec_list:
- spec_pair = (deepcopy(source_spec), deepcopy(target_spec))
- source_shard_list = self._convert_str_to_shard_list(source_spec)
- target_shard_list = self._convert_str_to_shard_list(target_spec)
+ source_shard_list = cls._convert_str_to_shard_list(source_spec)
+ target_shard_list = cls._convert_str_to_shard_list(target_spec)
# source same as target
if source_shard_list == target_shard_list:
@@ -112,30 +130,27 @@ class DimSpec:
else:
difference = NAN
- difference_dict[spec_pair] = difference
+ difference_dict[(source_spec, target_spec)] = difference
- self.difference_dict = difference_dict
+ return difference_dict
- def dim_diff(self, other):
+ @staticmethod
+ def _convert_str_to_shard_list(str_spec):
"""
- The difference between two _DimSpec.
+ Convert str_spec into shard_list.
Argument:
- other(_DimSpec): the dim spec to compare with.
-
- Return:
- difference(int): the difference between two _DimSpec.
-
- Example:
- dim_spec = _DimSpec([0])
- other_dim_spec = _DimSpec([0, 1])
- print(dim_spec.difference(other_dim_spec))
-
- Output:
- 5
+ str_spec(str): dim spec in str type.
"""
- difference = self.difference_dict[(str(self), str(other))]
- return difference
+
+ if str_spec == "R":
+ return []
+ if str_spec == "S0":
+ return [0]
+ if str_spec == "S1":
+ return [1]
+ if str_spec == "S01":
+ return [0, 1]
class ShardingSpec:
diff --git a/colossalai/tensor/sharding_spec.py b/colossalai/tensor/sharding_spec.py
index b78ef6d97..fb42afab7 100644
--- a/colossalai/tensor/sharding_spec.py
+++ b/colossalai/tensor/sharding_spec.py
@@ -1,5 +1,4 @@
import operator
-from copy import deepcopy
from functools import reduce
import torch
@@ -27,10 +26,11 @@ class _DimSpec:
Otherwise, the element in shard_list means the data will be sharded in that dimension.
"""
+ _DIFFERENCE_DICT = None
+
def __init__(self, shard_list):
self.is_replica = len(shard_list) == 0
self.shard_list = shard_list
- self.build_difference_2d_dict()
def __eq__(self, other):
return str(self) == str(other)
@@ -43,27 +43,46 @@ class _DimSpec:
target += str(dim)
return target
- def _convert_str_to_shard_list(self, str_spec):
+ @property
+ def difference_dict(self):
"""
- Convert str_spec into shard_list.
+ Returns the difference dict, and lazily initializes it when needed
+
+ Return:
+ difference_dict(Dict[Tuple[int, int], Union[int, float, str]]):
+ difference dict
+ """
+ if self._DIFFERENCE_DICT is None:
+ self._DIFFERENCE_DICT = self._build_difference_2d_dict()
+
+ return self._DIFFERENCE_DICT
+
+ def difference(self, other):
+ """
+ The difference between two _DimSpec.
Argument:
- str_spec(str): dim spec in str type.
+ other(_DimSpec): the dim spec to compare with.
+
+ Return:
+ difference(int): the difference between two _DimSpec.
+
+ Example:
+ dim_spec = _DimSpec([0])
+ other_dim_spec = _DimSpec([0, 1])
+ print(dim_spec.difference(other_dim_spec))
+
+ Output:
+ 5
"""
+ difference = self.difference_dict[(str(self), str(other))]
+ return difference
- if str_spec == "R":
- return []
- if str_spec == "S0":
- return [0]
- if str_spec == "S1":
- return [1]
- if str_spec == "S01":
- return [0, 1]
-
- def build_difference_2d_dict(self):
+ @classmethod
+ def _build_difference_2d_dict(cls):
"""
Build a difference mapping for 2D device mesh case. It will be used to
- compute the difference between DimSpec pairs.
+ compute the difference between _DimSpec pairs.
"""
source_spec_list = ["R", "S0", "S1", "S01"]
@@ -71,9 +90,8 @@ class _DimSpec:
difference_dict = {}
for source_spec in source_spec_list:
for target_spec in target_spec_list:
- spec_pair = (deepcopy(source_spec), deepcopy(target_spec))
- source_shard_list = self._convert_str_to_shard_list(source_spec)
- target_shard_list = self._convert_str_to_shard_list(target_spec)
+ source_shard_list = cls._convert_str_to_shard_list(source_spec)
+ target_shard_list = cls._convert_str_to_shard_list(target_spec)
# source same as target
if source_shard_list == target_shard_list:
@@ -116,30 +134,27 @@ class _DimSpec:
else:
difference = NAN
- difference_dict[spec_pair] = difference
+ difference_dict[(source_spec, target_spec)] = difference
- self.difference_dict = difference_dict
+ return difference_dict
- def difference(self, other):
+ @staticmethod
+ def _convert_str_to_shard_list(str_spec):
"""
- The difference between two _DimSpec.
+ Convert str_spec into shard_list.
Argument:
- other(_DimSpec): the dim spec to compare with.
-
- Return:
- difference(int): the difference between two _DimSpec.
-
- Example:
- dim_spec = _DimSpec([0])
- other_dim_spec = _DimSpec([0, 1])
- print(dim_spec.difference(other_dim_spec))
-
- Output:
- 5
+ str_spec(str): dim spec in str type.
"""
- difference = self.difference_dict[(str(self), str(other))]
- return difference
+
+ if str_spec == "R":
+ return []
+ if str_spec == "S0":
+ return [0]
+ if str_spec == "S1":
+ return [1]
+ if str_spec == "S01":
+ return [0, 1]
class ShardingSpecException(Exception):
diff --git a/colossalai/zero/low_level/bookkeeping/gradient_store.py b/colossalai/zero/low_level/bookkeeping/gradient_store.py
index e24a67f9d..8b6d403f1 100644
--- a/colossalai/zero/low_level/bookkeeping/gradient_store.py
+++ b/colossalai/zero/low_level/bookkeeping/gradient_store.py
@@ -19,7 +19,6 @@ class GradientStore(BaseStore):
"""
self._grads_of_params = dict()
# stage 2
- self._partition_grads = partition_grad
self._working_index = 0 if partition_grad else self._local_rank
# for zero2, it's `param_id: [grad_local_rank]`
self.grad_to_param_mapping = dict()
@@ -91,7 +90,7 @@ class GradientStore(BaseStore):
return grad_list
- def get_working_grad_by_param_id(self, param_id) -> Tensor:
+ def get_working_grad_by_param_id(self, param_id) -> Optional[Tensor]:
"""
Return the working gradient for the specified parameter.
@@ -112,6 +111,7 @@ class GradientStore(BaseStore):
def reset_all_gradients(self):
self._grads_of_params = dict()
+ self.grad_to_param_mapping = dict()
def get_param_id_for_grad(self, grad: Tensor) -> Optional[int]:
"""Return the id of a parameter which the gradient slice belongs to
diff --git a/colossalai/zero/low_level/low_level_optim.py b/colossalai/zero/low_level/low_level_optim.py
index 1353071c5..458e6e41a 100644
--- a/colossalai/zero/low_level/low_level_optim.py
+++ b/colossalai/zero/low_level/low_level_optim.py
@@ -21,9 +21,11 @@ from colossalai.amp.naive_amp.mixed_precision_mixin import (
from colossalai.interface import OptimizerWrapper
from colossalai.logging import get_dist_logger
from colossalai.quantization.fp8 import all_gather_into_tensor_flat_fp8, all_reduce_fp8, reduce_scatter_fp8
+from colossalai.tensor.moe_tensor.api import is_moe_tensor
from ._utils import calculate_global_norm_from_list, has_inf_or_nan, release_param_grad, sync_tensor
from .bookkeeping import BucketStore, GradientStore, TensorBucket
+from .zero_hook import set_all_gather_handle, wait_all_gather_handle
class LowLevelZeroFP16MixedPrecisionMixin(FP16MixedPrecisionMixin):
@@ -66,7 +68,7 @@ class LowLevelZeroOptimizer(OptimizerWrapper):
def __init__(
self,
optimizer: Optimizer,
- pg_to_param_list: Dict[ProcessGroup, List[nn.Parameter]] = None,
+ pg_to_param_list: Optional[Dict[ProcessGroup, List[nn.Parameter]]] = None,
initial_scale: int = 2**16, # grad scaler config
min_scale: int = 1,
growth_factor: float = 2.0,
@@ -84,6 +86,7 @@ class LowLevelZeroOptimizer(OptimizerWrapper):
dp_process_group: Optional[ProcessGroup] = None,
forced_dtype: Optional[torch.dtype] = None,
master_weights: bool = True, # master weights
+ overlap_allgather: bool = False,
fp8_communication: bool = False,
):
super(LowLevelZeroOptimizer, self).__init__(optim=optimizer)
@@ -92,7 +95,7 @@ class LowLevelZeroOptimizer(OptimizerWrapper):
self._logger = get_dist_logger()
self._verbose = verbose
- if dp_process_group is not None and pg_to_param_list is not None:
+ if (dp_process_group is not None) and (pg_to_param_list is not None):
raise ValueError("dp_process_group and pg_to_param_list should not be provided at the same time.")
if pg_to_param_list is None:
@@ -123,6 +126,7 @@ class LowLevelZeroOptimizer(OptimizerWrapper):
# communication params
self._overlap_communication = overlap_communication
+ self._overlap_allgather = overlap_allgather
self._reduce_bucket_size = reduce_bucket_size
self._communication_dtype = communication_dtype
self._fp8_communication = fp8_communication
@@ -148,6 +152,8 @@ class LowLevelZeroOptimizer(OptimizerWrapper):
# record the padding size of each param
self._padding_map = dict()
+ # padded working param is all-gather buffer and it shares the same memory with working param
+ self._working_param_to_padded_working_param = dict()
# mapping working param and master param
self.master_to_working_param = dict()
@@ -248,11 +254,12 @@ class LowLevelZeroOptimizer(OptimizerWrapper):
with torch.no_grad():
if padding_size > 0:
padding_param = torch.nn.functional.pad(param.data.view(-1), [0, padding_size])
- # reset working params' ptr when no master weights
- if self._master_weights == False:
- param.data = padding_param[: param.numel()].view(param.shape)
+ # # reset working params' ptr when no master weights
+ # if self._master_weights == False:
+ param.data = padding_param[: param.numel()].view(param.shape)
else:
padding_param = param.data.view(-1)
+ self._working_param_to_padded_working_param[param] = padding_param
splited_params = padding_param.split(
padding_param.numel() // self.pid_to_bucket_store[id(param)].world_size
@@ -261,7 +268,7 @@ class LowLevelZeroOptimizer(OptimizerWrapper):
# use fp32 when master_weights is True
if self._master_weights is True:
- splited_param_current_rank = splited_params.detach().float().to(device)
+ splited_param_current_rank = splited_params.detach().clone().float().to(device)
else:
splited_param_current_rank = splited_params
@@ -338,21 +345,21 @@ class LowLevelZeroOptimizer(OptimizerWrapper):
self._update_unpartitoned_grad(bucket_store, grad_in_bucket.values(), flat_grads_per_rank, group_id)
else:
flat_grads_list = list(flat_grads.split(len(flat_grads) // bucket_store.world_size))
- recieved_grad = torch.zeros_like(flat_grads_list[0])
+ received_grad = torch.zeros_like(flat_grads_list[0])
if self._fp8_communication:
reduce_scatter_fp8(
- recieved_grad,
+ received_grad,
flat_grads_list,
group=bucket_store.torch_pg,
)
else:
- dist.reduce_scatter(recieved_grad, flat_grads_list, group=bucket_store.torch_pg)
+ dist.reduce_scatter(received_grad, flat_grads_list, group=bucket_store.torch_pg)
- if recieved_grad.dtype != grad_dtype:
- recieved_grad = recieved_grad.to(grad_dtype)
+ if received_grad.dtype != grad_dtype:
+ received_grad = received_grad.to(grad_dtype)
grad_in_bucket_current_rank = bucket_store.get_grad()[bucket_store.local_rank]
- self._update_partitoned_grad(bucket_store, grad_in_bucket_current_rank, recieved_grad, group_id, 1)
+ self._update_partitoned_grad(bucket_store, grad_in_bucket_current_rank, received_grad, group_id, 1)
bucket_store.reset()
@@ -562,25 +569,29 @@ class LowLevelZeroOptimizer(OptimizerWrapper):
working_param = real_working_params[group_id][idx]
param_to_gather = master_param.to(device).to(self._dtype)
pg = self.param_to_pg[working_param]
- if param_to_gather.numel() > self.pg_to_tensor_bucket[pg].max_size:
- buffer_tensor = torch.empty_like(
- torch.cat([param_to_gather for _ in range(dist.get_world_size(pg))])
- )
- if self._fp8_communication:
- all_gather_into_tensor_flat_fp8(buffer_tensor, param_to_gather, pg, fp8_format="e4m3")
- else:
- dist.all_gather_into_tensor(buffer_tensor, param_to_gather, pg)
- working_param.data.copy_(buffer_tensor[: working_param.numel()].reshape_as(working_param))
- continue
- try:
- self.pg_to_tensor_bucket[pg].add_to_bucket(param_to_gather, write_back_tensor=working_param)
- except RuntimeError:
- self.pg_to_tensor_bucket[pg].all_gather(pg, fp8_communication=self._fp8_communication)
- self.pg_to_tensor_bucket[pg].add_to_bucket(param_to_gather, write_back_tensor=working_param)
+ padded_working_param = self._working_param_to_padded_working_param[working_param]
+ if self._overlap_allgather:
+ handle = dist.all_gather_into_tensor(padded_working_param, param_to_gather, pg, async_op=True)
+ set_all_gather_handle(working_param, handle)
+ else:
+ if param_to_gather.numel() > self.pg_to_tensor_bucket[pg].max_size:
+ if self._fp8_communication:
+ all_gather_into_tensor_flat_fp8(
+ padded_working_param, param_to_gather, pg, fp8_format="e4m3"
+ )
+ else:
+ dist.all_gather_into_tensor(padded_working_param, param_to_gather, pg)
+ continue
+ try:
+ self.pg_to_tensor_bucket[pg].add_to_bucket(param_to_gather, write_back_tensor=working_param)
+ except RuntimeError:
+ self.pg_to_tensor_bucket[pg].all_gather(pg, fp8_communication=self._fp8_communication)
+ self.pg_to_tensor_bucket[pg].add_to_bucket(param_to_gather, write_back_tensor=working_param)
self.optim.param_groups[group_id]["params"] = self._master_param_groups_of_current_rank[group_id]
- for pg, tensor_bucket in self.pg_to_tensor_bucket.items():
- if not tensor_bucket.is_empty():
- tensor_bucket.all_gather(pg, fp8_communication=self._fp8_communication)
+ if not self._overlap_allgather:
+ for pg, tensor_bucket in self.pg_to_tensor_bucket.items():
+ if not tensor_bucket.is_empty():
+ tensor_bucket.all_gather(pg, fp8_communication=self._fp8_communication)
def _compute_grad_norm(self, dp_pg: ProcessGroup, gradients: List[Tensor], norm_type: int = 2) -> float:
r"""
@@ -657,6 +668,11 @@ class LowLevelZeroOptimizer(OptimizerWrapper):
for group_id in range(self.num_param_groups):
param_group = self._working_param_groups[group_id]
for param in param_group:
+ if is_moe_tensor(param) and param.requires_grad and param.grad is None:
+ # TODO better of of doing this
+ # assign zero grad to unrouted expert to avoid hang during grad reduction
+ param.grad = torch.zeros_like(param)
+
if param.requires_grad and param.grad is not None:
self._add_to_bucket(param, group_id)
@@ -815,8 +831,8 @@ class LowLevelZeroOptimizer(OptimizerWrapper):
"""
for p in model.parameters():
p_id = id(p)
- pg = self.param_to_pg[p]
if p_id in self.working_to_master_param:
+ pg = self.param_to_pg[p]
master_param = self.working_to_master_param[p_id]
padding_size = self.get_param_padding_size(p)
working_param = p.data.view(-1)
@@ -877,13 +893,12 @@ class LowLevelZeroOptimizer(OptimizerWrapper):
def get_param_grad(self, working_param: nn.Parameter) -> Tensor:
grad_store = self.pid_to_grad_store[id(working_param)]
- partial_grad = grad_store.get_working_grad_by_param_id(id(working_param))
- if partial_grad is None:
+ grad = grad_store.get_working_grad_by_param_id(id(working_param))
+ if grad is None:
return None
- tensor_list = [torch.empty_like(partial_grad) for _ in range(grad_store.world_size)]
- dist.all_gather(tensor_list, partial_grad, group=grad_store.torch_pg)
- grad_flat = torch.cat(tensor_list, dim=0)
- return grad_flat[: working_param.numel()].reshape_as(working_param)
+ grad_flat = torch.empty((grad_store.world_size, *grad.shape), dtype=grad.dtype, device=grad.device)
+ dist.all_gather_into_tensor(grad_flat, grad, group=grad_store.torch_pg)
+ return grad_flat.view(-1)[: working_param.numel()].view_as(working_param)
def get_working_grads_by_group_id(self, group_id: int) -> List[Tensor]:
working_grads = []
@@ -908,3 +923,7 @@ class LowLevelZeroOptimizer(OptimizerWrapper):
def get_partitioned_gradients_by_param_id(self, group_id: int, param_id: int) -> List:
grad_store = self.pid_to_grad_store[param_id]
return grad_store.get_partitioned_gradients_by_param_id(group_id, param_id)
+
+ def _force_wait_all_gather(self):
+ for param in self._working_param_to_padded_working_param.keys():
+ wait_all_gather_handle(param)
diff --git a/colossalai/zero/low_level/zero_hook.py b/colossalai/zero/low_level/zero_hook.py
new file mode 100644
index 000000000..20f9ef31a
--- /dev/null
+++ b/colossalai/zero/low_level/zero_hook.py
@@ -0,0 +1,33 @@
+from typing import List
+
+from torch._tensor import Tensor
+
+from colossalai.tensor.param_op_hook import ColoParamOpHook
+
+_ALL_GATHER_HANDLE = "_all_gather_handle"
+
+
+def wait_all_gather_handle(p):
+ if hasattr(p, _ALL_GATHER_HANDLE):
+ handle = getattr(p, _ALL_GATHER_HANDLE)
+ handle.wait()
+ delattr(p, _ALL_GATHER_HANDLE)
+
+
+def set_all_gather_handle(p, handle):
+ setattr(p, _ALL_GATHER_HANDLE, handle)
+
+
+class ZeroOpHook(ColoParamOpHook):
+ def pre_forward(self, params: List[Tensor]) -> None:
+ for p in params:
+ wait_all_gather_handle(p)
+
+ def post_forward(self, params: List[Tensor]) -> None:
+ pass
+
+ def pre_backward(self, params: List[Tensor]) -> None:
+ pass
+
+ def post_backward(self, params: List[Tensor]) -> None:
+ pass
diff --git a/examples/inference/stable_diffusion/README.md b/examples/inference/stable_diffusion/README.md
new file mode 100644
index 000000000..c11b98043
--- /dev/null
+++ b/examples/inference/stable_diffusion/README.md
@@ -0,0 +1,22 @@
+## File Structure
+```
+|- sd3_generation.py: an example of how to use Colossalai Inference Engine to generate result by loading Diffusion Model.
+|- compute_metric.py: compare the quality of images w/o some acceleration method like Distrifusion
+|- benchmark_sd3.py: benchmark the performance of our InferenceEngine
+|- run_benchmark.sh: run benchmark command
+```
+Note: compute_metric.py need some dependencies which need `pip install -r requirements.txt`, `requirements.txt` is in `examples/inference/stable_diffusion/`
+
+## Run Inference
+
+The provided example `sd3_generation.py` is an example to configure, initialize the engine, and run inference on provided model. We've added `DiffusionPipeline` as model class, and the script is good to run inference with StableDiffusion 3.
+
+For a basic setting, you could run the example by:
+```bash
+colossalai run --nproc_per_node 1 sd3_generation.py -m PATH_MODEL -p "hello world"
+```
+
+Run multi-GPU inference (Patched Parallelism), as in the following example using 2 GPUs:
+```bash
+colossalai run --nproc_per_node 2 sd3_generation.py -m PATH_MODEL
+```
diff --git a/examples/inference/stable_diffusion/benchmark_sd3.py b/examples/inference/stable_diffusion/benchmark_sd3.py
new file mode 100644
index 000000000..19db57c33
--- /dev/null
+++ b/examples/inference/stable_diffusion/benchmark_sd3.py
@@ -0,0 +1,179 @@
+import argparse
+import json
+import time
+from contextlib import nullcontext
+
+import torch
+import torch.distributed as dist
+from diffusers import DiffusionPipeline
+
+import colossalai
+from colossalai.inference.config import DiffusionGenerationConfig, InferenceConfig
+from colossalai.inference.core.engine import InferenceEngine
+from colossalai.testing import clear_cache_before_run, rerun_if_address_is_in_use, spawn
+
+GIGABYTE = 1024**3
+MEGABYTE = 1024 * 1024
+
+_DTYPE_MAPPING = {
+ "fp16": torch.float16,
+ "bf16": torch.bfloat16,
+ "fp32": torch.float32,
+}
+
+
+def log_generation_time(log_data, log_file):
+ with open(log_file, "a") as f:
+ json.dump(log_data, f, indent=2)
+ f.write("\n")
+
+
+def warmup(engine, args):
+ for _ in range(args.n_warm_up_steps):
+ engine.generate(
+ prompts=["hello world"],
+ generation_config=DiffusionGenerationConfig(
+ num_inference_steps=args.num_inference_steps, height=args.height[0], width=args.width[0]
+ ),
+ )
+
+
+def profile_context(args):
+ return (
+ torch.profiler.profile(
+ record_shapes=True,
+ with_stack=True,
+ with_modules=True,
+ activities=[
+ torch.profiler.ProfilerActivity.CPU,
+ torch.profiler.ProfilerActivity.CUDA,
+ ],
+ )
+ if args.profile
+ else nullcontext()
+ )
+
+
+def log_and_profile(h, w, avg_time, log_msg, args, model_name, mode, prof=None):
+ log_data = {
+ "mode": mode,
+ "model": model_name,
+ "batch_size": args.batch_size,
+ "patched_parallel_size": args.patched_parallel_size,
+ "num_inference_steps": args.num_inference_steps,
+ "height": h,
+ "width": w,
+ "dtype": args.dtype,
+ "profile": args.profile,
+ "n_warm_up_steps": args.n_warm_up_steps,
+ "n_repeat_times": args.n_repeat_times,
+ "avg_generation_time": avg_time,
+ "log_message": log_msg,
+ }
+
+ if args.log:
+ log_file = f"examples/inference/stable_diffusion/benchmark_{model_name}_{mode}.json"
+ log_generation_time(log_data=log_data, log_file=log_file)
+
+ if args.profile:
+ file = f"examples/inference/stable_diffusion/benchmark_{model_name}_{mode}_prof.json"
+ prof.export_chrome_trace(file)
+
+
+def benchmark_colossalai(rank, world_size, port, args):
+ colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
+ from colossalai.cluster.dist_coordinator import DistCoordinator
+
+ coordinator = DistCoordinator()
+
+ inference_config = InferenceConfig(
+ dtype=args.dtype,
+ patched_parallelism_size=args.patched_parallel_size,
+ )
+ engine = InferenceEngine(args.model, inference_config=inference_config, verbose=False)
+
+ warmup(engine, args)
+
+ for h, w in zip(args.height, args.width):
+ with profile_context(args) as prof:
+ start = time.perf_counter()
+ for _ in range(args.n_repeat_times):
+ engine.generate(
+ prompts=["hello world"],
+ generation_config=DiffusionGenerationConfig(
+ num_inference_steps=args.num_inference_steps, height=h, width=w
+ ),
+ )
+ end = time.perf_counter()
+
+ avg_time = (end - start) / args.n_repeat_times
+ log_msg = f"[ColossalAI]avg generation time for h({h})xw({w}) is {avg_time:.2f}s"
+ coordinator.print_on_master(log_msg)
+
+ if dist.get_rank() == 0:
+ log_and_profile(h, w, avg_time, log_msg, args, args.model.split("/")[-1], "colossalai", prof=prof)
+
+
+def benchmark_diffusers(args):
+ model = DiffusionPipeline.from_pretrained(args.model, torch_dtype=_DTYPE_MAPPING[args.dtype]).to("cuda")
+
+ for _ in range(args.n_warm_up_steps):
+ model(
+ prompt="hello world",
+ num_inference_steps=args.num_inference_steps,
+ height=args.height[0],
+ width=args.width[0],
+ )
+
+ for h, w in zip(args.height, args.width):
+ with profile_context(args) as prof:
+ start = time.perf_counter()
+ for _ in range(args.n_repeat_times):
+ model(prompt="hello world", num_inference_steps=args.num_inference_steps, height=h, width=w)
+ end = time.perf_counter()
+
+ avg_time = (end - start) / args.n_repeat_times
+ log_msg = f"[Diffusers]avg generation time for h({h})xw({w}) is {avg_time:.2f}s"
+ print(log_msg)
+
+ log_and_profile(h, w, avg_time, log_msg, args, args.model.split("/")[-1], "diffusers", prof)
+
+
+@rerun_if_address_is_in_use()
+@clear_cache_before_run()
+def benchmark(args):
+ if args.mode == "colossalai":
+ spawn(benchmark_colossalai, nprocs=args.patched_parallel_size, args=args)
+ elif args.mode == "diffusers":
+ benchmark_diffusers(args)
+
+
+"""
+# enable log
+python examples/inference/stable_diffusion/benchmark_sd3.py -m "PixArt-alpha/PixArt-XL-2-1024-MS" -p 2 --mode colossalai --log
+python examples/inference/stable_diffusion/benchmark_sd3.py -m "PixArt-alpha/PixArt-XL-2-1024-MS" --mode diffusers --log
+
+# enable profiler
+python examples/inference/stable_diffusion/benchmark_sd3.py -m "stabilityai/stable-diffusion-3-medium-diffusers" -p 2 --mode colossalai --n_warm_up_steps 3 --n_repeat_times 1 --profile --num_inference_steps 20
+python examples/inference/stable_diffusion/benchmark_sd3.py -m "PixArt-alpha/PixArt-XL-2-1024-MS" -p 2 --mode colossalai --n_warm_up_steps 3 --n_repeat_times 1 --profile --num_inference_steps 20
+python examples/inference/stable_diffusion/benchmark_sd3.py -m "PixArt-alpha/PixArt-XL-2-1024-MS" --mode diffusers --n_warm_up_steps 3 --n_repeat_times 1 --profile --num_inference_steps 20
+"""
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("-b", "--batch_size", type=int, default=1, help="Batch size")
+ parser.add_argument("-p", "--patched_parallel_size", type=int, default=1, help="Patched Parallelism size")
+ parser.add_argument("-n", "--num_inference_steps", type=int, default=50, help="Number of inference steps")
+ parser.add_argument("-H", "--height", type=int, nargs="+", default=[1024, 2048], help="Height list")
+ parser.add_argument("-w", "--width", type=int, nargs="+", default=[1024, 2048], help="Width list")
+ parser.add_argument("--dtype", type=str, default="fp16", choices=["fp16", "fp32", "bf16"], help="Data type")
+ parser.add_argument("--n_warm_up_steps", type=int, default=3, help="Number of warm up steps")
+ parser.add_argument("--n_repeat_times", type=int, default=5, help="Number of repeat times")
+ parser.add_argument("--profile", default=False, action="store_true", help="Enable torch profiler")
+ parser.add_argument("--log", default=False, action="store_true", help="Enable logging")
+ parser.add_argument("-m", "--model", default="stabilityai/stable-diffusion-3-medium-diffusers", help="Model path")
+ parser.add_argument(
+ "--mode", default="colossalai", choices=["colossalai", "diffusers"], help="Inference framework mode"
+ )
+ args = parser.parse_args()
+ benchmark(args)
diff --git a/examples/inference/stable_diffusion/compute_metric.py b/examples/inference/stable_diffusion/compute_metric.py
new file mode 100644
index 000000000..14c92501b
--- /dev/null
+++ b/examples/inference/stable_diffusion/compute_metric.py
@@ -0,0 +1,80 @@
+# Code from https://github.com/mit-han-lab/distrifuser/blob/main/scripts/compute_metrics.py
+import argparse
+import os
+
+import numpy as np
+import torch
+from cleanfid import fid
+from PIL import Image
+from torch.utils.data import DataLoader, Dataset
+from torchmetrics.image import LearnedPerceptualImagePatchSimilarity, PeakSignalNoiseRatio
+from torchvision.transforms import Resize
+from tqdm import tqdm
+
+
+def read_image(path: str):
+ """
+ input: path
+ output: tensor (C, H, W)
+ """
+ img = np.asarray(Image.open(path))
+ if len(img.shape) == 2:
+ img = np.repeat(img[:, :, None], 3, axis=2)
+ img = torch.from_numpy(img).permute(2, 0, 1)
+ return img
+
+
+class MultiImageDataset(Dataset):
+ def __init__(self, root0, root1, is_gt=False):
+ super().__init__()
+ self.root0 = root0
+ self.root1 = root1
+ file_names0 = os.listdir(root0)
+ file_names1 = os.listdir(root1)
+
+ self.image_names0 = sorted([name for name in file_names0 if name.endswith(".png") or name.endswith(".jpg")])
+ self.image_names1 = sorted([name for name in file_names1 if name.endswith(".png") or name.endswith(".jpg")])
+ self.is_gt = is_gt
+ assert len(self.image_names0) == len(self.image_names1)
+
+ def __len__(self):
+ return len(self.image_names0)
+
+ def __getitem__(self, idx):
+ img0 = read_image(os.path.join(self.root0, self.image_names0[idx]))
+ if self.is_gt:
+ # resize to 1024 x 1024
+ img0 = Resize((1024, 1024))(img0)
+ img1 = read_image(os.path.join(self.root1, self.image_names1[idx]))
+
+ batch_list = [img0, img1]
+ return batch_list
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--batch_size", type=int, default=64)
+ parser.add_argument("--num_workers", type=int, default=8)
+ parser.add_argument("--is_gt", action="store_true")
+ parser.add_argument("--input_root0", type=str, required=True)
+ parser.add_argument("--input_root1", type=str, required=True)
+ args = parser.parse_args()
+
+ psnr = PeakSignalNoiseRatio(data_range=(0, 1), reduction="elementwise_mean", dim=(1, 2, 3)).to("cuda")
+ lpips = LearnedPerceptualImagePatchSimilarity(normalize=True).to("cuda")
+
+ dataset = MultiImageDataset(args.input_root0, args.input_root1, is_gt=args.is_gt)
+ dataloader = DataLoader(dataset, batch_size=args.batch_size, num_workers=args.num_workers)
+
+ progress_bar = tqdm(dataloader)
+ with torch.inference_mode():
+ for i, batch in enumerate(progress_bar):
+ batch = [img.to("cuda") / 255 for img in batch]
+ batch_size = batch[0].shape[0]
+ psnr.update(batch[0], batch[1])
+ lpips.update(batch[0], batch[1])
+ fid_score = fid.compute_fid(args.input_root0, args.input_root1)
+
+ print("PSNR:", psnr.compute().item())
+ print("LPIPS:", lpips.compute().item())
+ print("FID:", fid_score)
diff --git a/examples/inference/stable_diffusion/requirements.txt b/examples/inference/stable_diffusion/requirements.txt
new file mode 100644
index 000000000..c4e74162d
--- /dev/null
+++ b/examples/inference/stable_diffusion/requirements.txt
@@ -0,0 +1,3 @@
+torchvision
+torchmetrics
+cleanfid
diff --git a/examples/inference/stable_diffusion/run_benchmark.sh b/examples/inference/stable_diffusion/run_benchmark.sh
new file mode 100644
index 000000000..f3e45a335
--- /dev/null
+++ b/examples/inference/stable_diffusion/run_benchmark.sh
@@ -0,0 +1,42 @@
+#!/bin/bash
+
+models=("PixArt-alpha/PixArt-XL-2-1024-MS" "stabilityai/stable-diffusion-3-medium-diffusers")
+parallelism=(1 2 4 8)
+resolutions=(1024 2048 3840)
+modes=("colossalai" "diffusers")
+
+CUDA_VISIBLE_DEVICES_set_n_least_memory_usage() {
+ local n=${1:-"9999"}
+ echo "GPU Memory Usage:"
+ local FIRST_N_GPU_IDS=$(nvidia-smi --query-gpu=memory.used --format=csv \
+ | tail -n +2 \
+ | nl -v 0 \
+ | tee /dev/tty \
+ | sort -g -k 2 \
+ | awk '{print $1}' \
+ | head -n $n)
+ export CUDA_VISIBLE_DEVICES=$(echo $FIRST_N_GPU_IDS | sed 's/ /,/g')
+ echo "Now CUDA_VISIBLE_DEVICES is set to:"
+ echo "CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES"
+}
+
+for model in "${models[@]}"; do
+ for p in "${parallelism[@]}"; do
+ for resolution in "${resolutions[@]}"; do
+ for mode in "${modes[@]}"; do
+ if [[ "$mode" == "colossalai" && "$p" == 1 ]]; then
+ continue
+ fi
+ if [[ "$mode" == "diffusers" && "$p" != 1 ]]; then
+ continue
+ fi
+ CUDA_VISIBLE_DEVICES_set_n_least_memory_usage $p
+
+ cmd="python examples/inference/stable_diffusion/benchmark_sd3.py -m \"$model\" -p $p --mode $mode --log -H $resolution -w $resolution"
+
+ echo "Executing: $cmd"
+ eval $cmd
+ done
+ done
+ done
+done
diff --git a/examples/inference/stable_diffusion/sd3_generation.py b/examples/inference/stable_diffusion/sd3_generation.py
new file mode 100644
index 000000000..9e146c34b
--- /dev/null
+++ b/examples/inference/stable_diffusion/sd3_generation.py
@@ -0,0 +1,81 @@
+import argparse
+
+from diffusers import DiffusionPipeline
+from torch import bfloat16
+from torch import distributed as dist
+from torch import float16, float32
+
+import colossalai
+from colossalai.cluster import DistCoordinator
+from colossalai.inference.config import DiffusionGenerationConfig, InferenceConfig
+from colossalai.inference.core.engine import InferenceEngine
+
+# For Stable Diffusion 3, we'll use the following configuration
+MODEL_CLS = DiffusionPipeline
+
+TORCH_DTYPE_MAP = {
+ "fp16": float16,
+ "fp32": float32,
+ "bf16": bfloat16,
+}
+
+
+def infer(args):
+ # ==============================
+ # Launch colossalai, setup distributed environment
+ # ==============================
+ colossalai.launch_from_torch()
+ coordinator = DistCoordinator()
+
+ # ==============================
+ # Load model and tokenizer
+ # ==============================
+ model_path_or_name = args.model
+ model = MODEL_CLS.from_pretrained(model_path_or_name, torch_dtype=TORCH_DTYPE_MAP.get(args.dtype, None))
+
+ # ==============================
+ # Initialize InferenceEngine
+ # ==============================
+ coordinator.print_on_master(f"Initializing Inference Engine...")
+ inference_config = InferenceConfig(
+ dtype=args.dtype,
+ max_batch_size=args.max_batch_size,
+ tp_size=args.tp_size,
+ use_cuda_kernel=args.use_cuda_kernel,
+ patched_parallelism_size=dist.get_world_size(),
+ )
+ engine = InferenceEngine(model, inference_config=inference_config, verbose=True)
+
+ # ==============================
+ # Generation
+ # ==============================
+ coordinator.print_on_master(f"Generating...")
+ out = engine.generate(prompts=[args.prompt], generation_config=DiffusionGenerationConfig())[0]
+ if dist.get_rank() == 0:
+ out.save(f"cat_parallel_size{dist.get_world_size()}.jpg")
+ coordinator.print_on_master(out)
+
+
+# colossalai run --nproc_per_node 1 examples/inference/stable_diffusion/sd3_generation.py -m MODEL_PATH
+
+# colossalai run --nproc_per_node 1 examples/inference/stable_diffusion/sd3_generation.py -m "stabilityai/stable-diffusion-3-medium-diffusers" --tp_size 1
+# colossalai run --nproc_per_node 2 examples/inference/stable_diffusion/sd3_generation.py -m "stabilityai/stable-diffusion-3-medium-diffusers" --tp_size 1
+
+# colossalai run --nproc_per_node 1 examples/inference/stable_diffusion/sd3_generation.py -m "PixArt-alpha/PixArt-XL-2-1024-MS" --tp_size 1
+# colossalai run --nproc_per_node 2 examples/inference/stable_diffusion/sd3_generation.py -m "PixArt-alpha/PixArt-XL-2-1024-MS" --tp_size 1
+
+
+if __name__ == "__main__":
+ # ==============================
+ # Parse Arguments
+ # ==============================
+ parser = argparse.ArgumentParser()
+ parser.add_argument("-m", "--model", type=str, help="Path to the model or model name")
+ parser.add_argument("-t", "--tp_size", type=int, default=1, help="Tensor Parallelism size")
+ parser.add_argument("-p", "--prompt", type=str, default="A cat holding a sign that says hello world", help="Prompt")
+ parser.add_argument("-b", "--max_batch_size", type=int, default=1, help="Max batch size")
+ parser.add_argument("-d", "--dtype", type=str, default="fp16", help="Data type", choices=["fp16", "fp32", "bf16"])
+ parser.add_argument("--use_cuda_kernel", action="store_true", help="Use CUDA kernel, use Triton by default")
+ args = parser.parse_args()
+
+ infer(args)
diff --git a/examples/inference/stable_diffusion/test_ci.sh b/examples/inference/stable_diffusion/test_ci.sh
new file mode 100644
index 000000000..d0189431c
--- /dev/null
+++ b/examples/inference/stable_diffusion/test_ci.sh
@@ -0,0 +1,2 @@
+#!/bin/bash
+echo "Skip the test (this test is slow)"
diff --git a/examples/language/gpt/hybridparallelism/finetune.py b/examples/language/gpt/hybridparallelism/finetune.py
index 9b3a10160..f447adf69 100644
--- a/examples/language/gpt/hybridparallelism/finetune.py
+++ b/examples/language/gpt/hybridparallelism/finetune.py
@@ -1,4 +1,5 @@
import argparse
+from contextlib import nullcontext
from typing import Callable, List, Union
import evaluate
@@ -17,6 +18,7 @@ from colossalai.accelerator import get_accelerator
from colossalai.booster import Booster
from colossalai.booster.plugin import GeminiPlugin, HybridParallelPlugin, LowLevelZeroPlugin, TorchDDPPlugin
from colossalai.cluster import DistCoordinator
+from colossalai.lazy import LazyInitContext
from colossalai.nn.optimizer import HybridAdam
# ==============================
@@ -252,10 +254,16 @@ def main():
pad_token_id=data_builder.tokenizer.pad_token_id,
)
- if model_name == "gpt2":
- model = GPT2ForSequenceClassification.from_pretrained(model_name, config=cfg).cuda()
- else:
- raise RuntimeError
+ init_ctx = (
+ LazyInitContext(default_device=get_accelerator().get_current_device())
+ if isinstance(plugin, (GeminiPlugin))
+ else nullcontext()
+ )
+ with init_ctx:
+ if model_name == "gpt2":
+ model = GPT2ForSequenceClassification.from_pretrained(model_name, config=cfg).cuda()
+ else:
+ raise RuntimeError
# optimizer
no_decay = ["bias", "LayerNorm.weight"]
diff --git a/examples/language/llama/benchmark.py b/examples/language/llama/benchmark.py
index 8a35db1f7..e530e2d6a 100644
--- a/examples/language/llama/benchmark.py
+++ b/examples/language/llama/benchmark.py
@@ -98,6 +98,7 @@ def main():
parser.add_argument("--disable-async-reduce", action="store_true", help="Disable the asynchronous reduce operation")
parser.add_argument("--prefetch_num", type=int, default=0, help="chunk prefetch max number")
parser.add_argument("--no_cache", action="store_true")
+ parser.add_argument("--overlap_allgather", action="store_true")
args = parser.parse_args()
colossalai.launch_from_torch()
@@ -199,9 +200,9 @@ def main():
enable_flash_attention=args.xformers,
microbatch_size=args.mbs,
precision="bf16",
- dp_outside=False,
overlap_p2p=args.overlap,
enable_metadata_cache=not args.no_cache,
+ overlap_allgather=args.overlap_allgather,
**hybrid_kwargs,
)
elif args.plugin == "3d_cpu":
@@ -292,7 +293,7 @@ def main():
with get_profile_context(
args.profile,
args.ignore_steps,
- len(dataloader) - 1,
+ 1, # avoid creating massive log files
save_dir=f"profile/{time.strftime('%H:%M', time.localtime())}-{args.plugin}-llama-{args.config}",
) as prof:
if isinstance(plugin, HybridParallelPlugin) and args.pp > 1:
diff --git a/examples/language/opt/opt_benchmark.py b/examples/language/opt/opt_benchmark.py
index c2883d96c..ca9b63d1a 100755
--- a/examples/language/opt/opt_benchmark.py
+++ b/examples/language/opt/opt_benchmark.py
@@ -1,4 +1,5 @@
import time
+from contextlib import nullcontext
import torch
import tqdm
@@ -8,9 +9,11 @@ from transformers import AutoConfig, OPTForCausalLM
from transformers.utils.versions import require_version
import colossalai
+from colossalai.accelerator import get_accelerator
from colossalai.booster import Booster
from colossalai.booster.plugin import GeminiPlugin, LowLevelZeroPlugin, TorchDDPPlugin
from colossalai.cluster import DistCoordinator
+from colossalai.lazy import LazyInitContext
from colossalai.logging import disable_existing_loggers, get_dist_logger
from colossalai.nn.optimizer import HybridAdam
@@ -62,14 +65,6 @@ def main():
if args.mem_cap > 0:
colo_memory_cap(args.mem_cap)
- # Build OPT model
- config = AutoConfig.from_pretrained(args.model_name_or_path)
- model = OPTForCausalLM(config=config)
- logger.info(f"Finish loading model from {args.model_name_or_path}", ranks=[0])
-
- # Enable gradient checkpointing
- model.gradient_checkpointing_enable()
-
# Set plugin
booster_kwargs = {}
if args.plugin == "torch_ddp_fp16":
@@ -82,6 +77,19 @@ def main():
plugin = LowLevelZeroPlugin(initial_scale=2**5)
logger.info(f"Set plugin as {args.plugin}", ranks=[0])
+ # Build OPT model
+ init_ctx = (
+ LazyInitContext(default_device=get_accelerator().get_current_device())
+ if isinstance(plugin, (GeminiPlugin))
+ else nullcontext()
+ )
+ config = AutoConfig.from_pretrained(args.model_name_or_path)
+ with init_ctx:
+ model = OPTForCausalLM(config=config)
+ logger.info(f"Finish loading model from {args.model_name_or_path}", ranks=[0])
+
+ # Enable gradient checkpointing
+ model.gradient_checkpointing_enable()
# Set optimizer
optimizer = HybridAdam(model.parameters(), lr=args.learning_rate)
diff --git a/examples/language/opt/opt_train_demo.py b/examples/language/opt/opt_train_demo.py
index b5b50305c..50dfc7bff 100644
--- a/examples/language/opt/opt_train_demo.py
+++ b/examples/language/opt/opt_train_demo.py
@@ -1,3 +1,5 @@
+from contextlib import nullcontext
+
import datasets
import torch
import transformers
@@ -8,9 +10,11 @@ from transformers import AutoConfig, AutoTokenizer, OPTForCausalLM, get_linear_s
from transformers.utils.versions import require_version
import colossalai
+from colossalai.accelerator import get_accelerator
from colossalai.booster import Booster
from colossalai.booster.plugin import GeminiPlugin, HybridParallelPlugin, LowLevelZeroPlugin, TorchDDPPlugin
from colossalai.cluster import DistCoordinator
+from colossalai.lazy import LazyInitContext
from colossalai.logging import disable_existing_loggers, get_dist_logger
from colossalai.nn.optimizer import HybridAdam
@@ -78,14 +82,6 @@ def main():
datasets.utils.logging.set_verbosity_error()
transformers.utils.logging.set_verbosity_error()
- # Build OPT model
- config = AutoConfig.from_pretrained(args.model_name_or_path)
- model = OPTForCausalLM.from_pretrained(args.model_name_or_path, config=config)
- logger.info(f"Finish loading model from {args.model_name_or_path}", ranks=[0])
-
- # Enable gradient checkpointing
- model.gradient_checkpointing_enable()
-
# Set plugin
booster_kwargs = {}
if args.plugin == "torch_ddp_fp16":
@@ -110,6 +106,21 @@ def main():
logger.info(f"Set plugin as {args.plugin}", ranks=[0])
+ # Build OPT model
+ config = AutoConfig.from_pretrained(args.model_name_or_path)
+ # Build OPT model
+ init_ctx = (
+ LazyInitContext(default_device=get_accelerator().get_current_device())
+ if isinstance(plugin, (GeminiPlugin, HybridParallelPlugin))
+ else nullcontext()
+ )
+ with init_ctx:
+ model = OPTForCausalLM.from_pretrained(args.model_name_or_path, config=config)
+ logger.info(f"Finish loading model from {args.model_name_or_path}", ranks=[0])
+
+ # Enable gradient checkpointing
+ model.gradient_checkpointing_enable()
+
# Prepare tokenizer and dataloader
tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
dataset = NetflixDataset(tokenizer)
diff --git a/examples/language/performance_evaluator.py b/examples/language/performance_evaluator.py
index 6b8daf37d..ca4a02cd2 100644
--- a/examples/language/performance_evaluator.py
+++ b/examples/language/performance_evaluator.py
@@ -113,13 +113,13 @@ class PerformanceEvaluator:
self.disable = self.ignore_steps > 0 and step < self.ignore_steps
if self.disable:
return
- get_accelerator().synchronize()
+ # get_accelerator().synchronize()
self.timer.start()
def on_step_end(self, input_ids: Tensor, **kwargs) -> None:
if self.disable:
return
- get_accelerator().synchronize()
+ # get_accelerator().synchronize()
self.timer.end()
batch_size, seq_len = input_ids.shape
diff --git a/requirements/requirements-test.txt b/requirements/requirements-test.txt
index e4affc7f5..93a3690fe 100644
--- a/requirements/requirements-test.txt
+++ b/requirements/requirements-test.txt
@@ -1,4 +1,3 @@
-diffusers
pytest
coverage==7.2.3
git+https://github.com/hpcaitech/pytest-testmon
diff --git a/requirements/requirements.txt b/requirements/requirements.txt
index 27bbc3769..651eb66e8 100644
--- a/requirements/requirements.txt
+++ b/requirements/requirements.txt
@@ -8,7 +8,7 @@ click
fabric
contexttimer
ninja
-torch>=2.1.0,<2.3.0
+torch>=2.1.0,<=2.3.0
safetensors
einops
pydantic
@@ -23,3 +23,4 @@ rpyc==6.0.0
fastapi
uvicorn==0.29.0
galore_torch
+diffusers==0.29.0
diff --git a/tests/kit/model_zoo/transformers/__init__.py b/tests/kit/model_zoo/transformers/__init__.py
index 05c17f562..4adc38619 100644
--- a/tests/kit/model_zoo/transformers/__init__.py
+++ b/tests/kit/model_zoo/transformers/__init__.py
@@ -3,28 +3,17 @@ from .bert import *
from .blip2 import *
from .bloom import *
from .chatglm2 import *
+from .command import *
+from .deepseek import *
from .falcon import *
from .gpt import *
from .gptj import *
from .llama import *
+from .mistral import *
+from .mixtral import *
from .opt import *
+from .qwen2 import *
from .sam import *
from .t5 import *
from .vit import *
from .whisper import *
-
-try:
- from .mistral import *
-except ImportError:
- print("This version of transformers doesn't support mistral.")
-
-try:
- from .qwen2 import *
-except ImportError:
- print("This version of transformers doesn't support qwen2.")
-
-
-try:
- from .command import *
-except ImportError:
- print("This version of transformers doesn't support Command-R.")
diff --git a/tests/kit/model_zoo/transformers/deepseek.py b/tests/kit/model_zoo/transformers/deepseek.py
new file mode 100644
index 000000000..ad73640a5
--- /dev/null
+++ b/tests/kit/model_zoo/transformers/deepseek.py
@@ -0,0 +1,83 @@
+# modified from tests/kit/model_zoo/transformers/mistral.py
+import torch
+import transformers
+from transformers import AutoConfig
+
+from ..registry import ModelAttribute, model_zoo
+
+# ===============================
+# Register single-sentence Mixtral
+# ===============================
+
+
+def data_gen():
+ # Generated from following code snippet
+ #
+ # from transformers import AutoModelForCausalLM, AutoTokenizer
+ # tokenizer = AutoTokenizer.from_pretrained("mixtralai/Mixtral-7B-v0.1")
+ # input = 'My favourite condiment is vinegar' (last two words repeated to satisfy length requirement)
+ # tokenized_input = tokenizer([input], return_tensors="pt")
+ # input_ids = tokenized_input['input_ids']
+ # attention_mask = tokenized_input['attention_mask']
+ input_ids = torch.tensor([[1, 22, 55, 77, 532, 349, 43, 22]], dtype=torch.int64)
+ attention_mask = torch.tensor([[1, 1, 1, 1, 1, 1, 1, 1]], dtype=torch.int64)
+ return dict(input_ids=input_ids, attention_mask=attention_mask)
+
+
+def data_gen_for_lm():
+ # LM data gen
+ # the `labels` of LM is the token of the output, cause no padding, use `input_ids` as `labels`
+ data = data_gen()
+ data["labels"] = data["input_ids"].clone()
+ return data
+
+
+def data_gen_for_sequence_classification():
+ # sequence classification data gen
+ data = data_gen()
+ data["labels"] = torch.tensor([1], dtype=torch.int64)
+ return data
+
+
+# define output transform function
+output_transform_fn = lambda x: x
+
+# define loss function
+loss_fn_for_mixtral_model = lambda x: x[0].mean()
+loss_fn = lambda x: x.loss
+loss_fn_for_seq_classification = lambda output: output.logits.mean()
+
+
+def init_deepseek():
+
+ config = AutoConfig.from_pretrained(
+ "deepseek-ai/deepseek-moe-16b-base",
+ hidden_size=32,
+ intermediate_size=32,
+ moe_intermediate_size=32,
+ num_hidden_layers=2,
+ num_attention_heads=8,
+ num_key_value_heads=8,
+ # vocab_size=2200,
+ first_k_dense_replace=1,
+ attn_implementation="flash_attention_2",
+ torch_dtype="float16",
+ n_routed_experts=8,
+ trust_remote_code=True,
+ )
+
+ if hasattr(config, "pad_token_id"):
+ config.pad_token_id = config.eos_token_id
+ model = transformers.AutoModel.from_config(config, trust_remote_code=True)
+
+ return model
+
+
+model_zoo.register(
+ name="transformers_deepseek",
+ model_fn=init_deepseek,
+ data_gen_fn=data_gen,
+ output_transform_fn=output_transform_fn,
+ loss_fn=loss_fn_for_mixtral_model,
+ model_attribute=ModelAttribute(has_control_flow=True),
+)
diff --git a/tests/kit/model_zoo/transformers/mixtral.py b/tests/kit/model_zoo/transformers/mixtral.py
new file mode 100644
index 000000000..73c0e9e2c
--- /dev/null
+++ b/tests/kit/model_zoo/transformers/mixtral.py
@@ -0,0 +1,85 @@
+# modified from tests/kit/model_zoo/transformers/mistral.py
+import torch
+import transformers
+from transformers import MixtralConfig
+
+from ..registry import ModelAttribute, model_zoo
+
+# ===============================
+# Register single-sentence Mixtral
+# ===============================
+
+
+def data_gen():
+ # Generated from following code snippet
+ #
+ # from transformers import AutoModelForCausalLM, AutoTokenizer
+ # tokenizer = AutoTokenizer.from_pretrained("mixtralai/Mixtral-7B-v0.1")
+ # input = 'My favourite condiment is vinegar' (last two words repeated to satisfy length requirement)
+ # tokenized_input = tokenizer([input], return_tensors="pt")
+ # input_ids = tokenized_input['input_ids']
+ # attention_mask = tokenized_input['attention_mask']
+ input_ids = torch.tensor([[1, 22, 55, 77, 532, 349, 43, 22]], dtype=torch.int64)
+ attention_mask = torch.tensor([[1, 1, 1, 1, 1, 1, 1, 1]], dtype=torch.int64)
+ return dict(input_ids=input_ids, attention_mask=attention_mask)
+
+
+def data_gen_for_lm():
+ # LM data gen
+ # the `labels` of LM is the token of the output, cause no padding, use `input_ids` as `labels`
+ data = data_gen()
+ data["labels"] = data["input_ids"].clone()
+ return data
+
+
+def data_gen_for_sequence_classification():
+ # sequence classification data gen
+ data = data_gen()
+ data["labels"] = torch.tensor([1], dtype=torch.int64)
+ return data
+
+
+# define output transform function
+output_transform_fn = lambda x: x
+
+# define loss function
+loss_fn_for_mixtral_model = lambda x: x[0].mean()
+loss_fn = lambda x: x.loss
+loss_fn_for_seq_classification = lambda output: output.logits.mean()
+
+config = MixtralConfig(
+ hidden_size=32,
+ intermediate_size=32,
+ num_attention_heads=8,
+ num_hidden_layers=2,
+ vocab_size=1000,
+ output_router_logits=True,
+)
+
+if hasattr(config, "pad_token_id"):
+ config.pad_token_id = config.eos_token_id
+
+model_zoo.register(
+ name="transformers_mixtral",
+ model_fn=lambda: transformers.MixtralModel(config),
+ data_gen_fn=data_gen,
+ output_transform_fn=output_transform_fn,
+ loss_fn=loss_fn_for_mixtral_model,
+ model_attribute=ModelAttribute(has_control_flow=True),
+)
+# model_zoo.register(
+# name="transformers_mixtral_for_casual_lm",
+# model_fn=lambda: transformers.MixtralForCausalLM(config),
+# data_gen_fn=data_gen_for_lm,
+# output_transform_fn=output_transform_fn,
+# loss_fn=loss_fn,
+# model_attribute=ModelAttribute(has_control_flow=True),
+# )
+# model_zoo.register(
+# name="transformers_mixtral_for_sequence_classification",
+# model_fn=lambda: transformers.MixtralForSequenceClassification(config),
+# data_gen_fn=data_gen_for_sequence_classification,
+# output_transform_fn=output_transform_fn,
+# loss_fn=loss_fn_for_seq_classification,
+# model_attribute=ModelAttribute(has_control_flow=True),
+# )
diff --git a/tests/test_legacy/test_moe/moe_utils.py b/tests/test_legacy/test_moe/moe_utils.py
new file mode 100644
index 000000000..8c133849b
--- /dev/null
+++ b/tests/test_legacy/test_moe/moe_utils.py
@@ -0,0 +1,136 @@
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from torch.distributed import ProcessGroup
+
+from colossalai.booster.plugin.low_level_zero_plugin import LowLevelZeroModel
+from colossalai.legacy.engine.gradient_handler._base_gradient_handler import BaseGradientHandler
+from colossalai.legacy.engine.gradient_handler.utils import bucket_allreduce
+from colossalai.legacy.moe.manager import MOE_MANAGER
+from colossalai.legacy.moe.utils import get_moe_epsize_param_dict
+from colossalai.legacy.registry import GRADIENT_HANDLER
+from colossalai.tensor.moe_tensor.api import get_ep_group, get_ep_size, set_moe_tensor_ep_group
+
+
+def delete_moe_info(model):
+ for _, param in model.named_parameters():
+ if hasattr(param, "ep_group"):
+ delattr(param, "ep_group")
+
+
+class MoeModel(nn.Module):
+ def __init__(self, ep_group: ProcessGroup = None):
+ super().__init__()
+ self.test_embed = nn.Linear(4, 16, bias=False)
+ self.w1 = torch.nn.Parameter(torch.randn(16, 8))
+ if ep_group:
+ set_moe_tensor_ep_group(self.w1, ep_group)
+
+ def forward(self, x):
+ x = self.test_embed(x)
+ x = torch.matmul(x, self.w1)
+
+ return x
+
+
+@GRADIENT_HANDLER.register_module
+class MoeGradientHandler(BaseGradientHandler):
+ """A helper class to handle all-reduce operations in a data parallel group and
+ moe model parallel. A all-reduce collective communication will be operated in
+ :func:`handle_gradient` among a data parallel group.
+ For better performance, it bucketizes the gradients of all parameters that are
+ the same type to improve the efficiency of communication.
+
+ Args:
+ model (Module): Model where the gradients accumulate.
+ optimizer (Optimizer): Optimizer for updating the parameters.
+ """
+
+ def __init__(self, model, optimizer=None):
+ super().__init__(model, optimizer)
+
+ def handle_gradient(self):
+ """A method running an all-reduce operation in a data parallel group.
+ Then running an all-reduce operation for all parameters in experts
+ across moe model parallel group
+ """
+ if dist.get_world_size() > 1:
+ epsize_param_dict = get_moe_epsize_param_dict(self._model)
+
+ # epsize is 1, indicating the params are replicated among processes in data parallelism
+ # use the ParallelMode.DATA to get data parallel group
+ # reduce gradients for all parameters in data parallelism
+ if 1 in epsize_param_dict:
+ bucket_allreduce(param_list=epsize_param_dict[1])
+
+ for ep_size in epsize_param_dict:
+ if ep_size != 1 and ep_size != MOE_MANAGER.world_size:
+ bucket_allreduce(
+ param_list=epsize_param_dict[ep_size], group=MOE_MANAGER.parallel_info_dict[ep_size].dp_group
+ )
+
+
+def assert_not_equal_in_group(tensor, process_group=None):
+ # all gather tensors from different ranks
+ world_size = dist.get_world_size(process_group)
+ tensor_list = [torch.empty_like(tensor) for _ in range(world_size)]
+ dist.all_gather(tensor_list, tensor, group=process_group)
+
+ # check if they are equal one by one
+ for i in range(world_size - 1):
+ a = tensor_list[i]
+ b = tensor_list[i + 1]
+ assert not torch.allclose(a, b), (
+ f"expected tensors on rank {i} and {i + 1} not to be equal " f"but they are, {a} vs {b}"
+ )
+
+
+def run_fwd_bwd(model, data, label, criterion, optimizer, enable_autocast=False):
+ model.train()
+ with torch.cuda.amp.autocast(enabled=enable_autocast):
+ if criterion:
+ y = model(data)
+ loss = criterion(y, label)
+ else:
+ loss = model(data, label)
+ loss = loss.float()
+
+ if isinstance(model, LowLevelZeroModel):
+ optimizer.backward(loss)
+ else:
+ loss.backward()
+ return y
+
+
+def sync_local_from_ep(local_model, ep_model, assert_grad_flag: bool = False) -> None:
+ """Sync the parameters of tp model from ep model
+
+ Args:
+ local_model (MoeModule)
+ ep_model (MoeModule)
+ """
+ for (local_name, local_param), (ep_name, ep_param) in zip(
+ local_model.named_parameters(), ep_model.named_parameters()
+ ):
+ if "experts" not in local_name:
+ if assert_grad_flag:
+ assert torch.allclose(local_param, ep_param), f"local_param: {local_param}, ep_param: {ep_param}"
+ assert torch.allclose(local_param.grad, ep_param.grad)
+ else:
+ local_param.data.copy_(ep_param.data)
+ continue
+
+ # gather param from ep model
+ param_list = [torch.zeros_like(ep_param) for _ in range(get_ep_size(ep_param))]
+ dist.all_gather(param_list, ep_param, group=get_ep_group(ep_param))
+ all_param = torch.cat(param_list, dim=0)
+ if assert_grad_flag:
+ grad_list = [torch.zeros_like(ep_param) for _ in range(get_ep_size(ep_param))]
+ dist.all_gather(grad_list, ep_param.grad, group=get_ep_group(ep_param))
+ all_grad = torch.cat(grad_list, dim=0)
+
+ if assert_grad_flag:
+ assert torch.allclose(local_param, all_param)
+ assert torch.allclose(local_param.grad, all_grad)
+ else:
+ local_param.data.copy_(all_param.data)
diff --git a/tests/test_moe/test_grad_handler.py b/tests/test_legacy/test_moe/test_grad_handler.py
similarity index 98%
rename from tests/test_moe/test_grad_handler.py
rename to tests/test_legacy/test_moe/test_grad_handler.py
index 25e61b091..3a782a6dd 100644
--- a/tests/test_moe/test_grad_handler.py
+++ b/tests/test_legacy/test_moe/test_grad_handler.py
@@ -5,7 +5,7 @@ import torch.nn as nn
import colossalai
from colossalai.accelerator import get_accelerator
-from colossalai.moe.manager import MOE_MANAGER
+from colossalai.legacy.moe.manager import MOE_MANAGER
# from colossalai.shardformer.layer.moe.layers import SparseMLP
from colossalai.testing import assert_equal_in_group, rerun_if_address_is_in_use, spawn
diff --git a/tests/test_moe/test_moe_group.py b/tests/test_legacy/test_moe/test_moe_group.py
similarity index 95%
rename from tests/test_moe/test_moe_group.py
rename to tests/test_legacy/test_moe/test_moe_group.py
index 89baf1d37..68dac4828 100644
--- a/tests/test_moe/test_moe_group.py
+++ b/tests/test_legacy/test_moe/test_moe_group.py
@@ -4,8 +4,8 @@ import torch.nn as nn
import colossalai
from colossalai.accelerator import get_accelerator
-from colossalai.moe.manager import MOE_MANAGER
-from colossalai.moe.utils import sync_moe_model_param
+from colossalai.legacy.moe.manager import MOE_MANAGER
+from colossalai.legacy.moe.utils import sync_moe_model_param
# from colossalai.shardformer.layer.moe import MLPExperts
from colossalai.testing import assert_equal_in_group, rerun_if_address_is_in_use, spawn
diff --git a/tests/test_moe/test_moe_hybrid_zero.py b/tests/test_legacy/test_moe/test_moe_hybrid_zero.py
similarity index 98%
rename from tests/test_moe/test_moe_hybrid_zero.py
rename to tests/test_legacy/test_moe/test_moe_hybrid_zero.py
index 513c4ebda..fdd6d956e 100644
--- a/tests/test_moe/test_moe_hybrid_zero.py
+++ b/tests/test_legacy/test_moe/test_moe_hybrid_zero.py
@@ -6,7 +6,7 @@ import colossalai
from colossalai.booster import Booster
from colossalai.booster.plugin import LowLevelZeroPlugin
from colossalai.booster.plugin.low_level_zero_plugin import LowLevelZeroModel
-from colossalai.moe.manager import MOE_MANAGER
+from colossalai.legacy.moe.manager import MOE_MANAGER
from colossalai.tensor.moe_tensor.api import is_moe_tensor
from colossalai.testing import rerun_if_address_is_in_use, spawn
from tests.test_moe.moe_utils import MoeModel
diff --git a/tests/test_moe/test_moe_load_balance.py b/tests/test_legacy/test_moe/test_moe_load_balance.py
similarity index 99%
rename from tests/test_moe/test_moe_load_balance.py
rename to tests/test_legacy/test_moe/test_moe_load_balance.py
index ddd3ea368..adf2dbc1c 100644
--- a/tests/test_moe/test_moe_load_balance.py
+++ b/tests/test_legacy/test_moe/test_moe_load_balance.py
@@ -6,7 +6,7 @@ import colossalai
from colossalai.booster import Booster
from colossalai.booster.plugin import LowLevelZeroPlugin
from colossalai.booster.plugin.low_level_zero_plugin import LowLevelZeroModel
-from colossalai.moe.manager import MOE_MANAGER
+from colossalai.legacy.moe.manager import MOE_MANAGER
# from colossalai.shardformer.layer.moe import apply_load_balance
from colossalai.tensor.moe_tensor.api import is_moe_tensor
diff --git a/tests/test_lora/test_lora.py b/tests/test_lora/test_lora.py
index b8daf775d..1ae17025d 100644
--- a/tests/test_lora/test_lora.py
+++ b/tests/test_lora/test_lora.py
@@ -9,7 +9,8 @@ from torch.optim import AdamW
import colossalai
from colossalai.booster import Booster
-from colossalai.booster.plugin import LowLevelZeroPlugin, TorchDDPPlugin
+from colossalai.booster.plugin import HybridParallelPlugin, LowLevelZeroPlugin, TorchDDPPlugin
+from colossalai.booster.plugin.hybrid_parallel_plugin import HybridParallelModule
from colossalai.testing import check_state_dict_equal, clear_cache_before_run, rerun_if_address_is_in_use, spawn
from tests.kit.model_zoo import model_zoo
from tests.test_checkpoint_io.utils import shared_tempdir
@@ -20,7 +21,7 @@ def check_fwd_bwd(model_fn, data_gen_fn, output_transform_fn, loss_fn, task_type
model = model_fn()
lora_config = LoraConfig(task_type=task_type, r=8, lora_alpha=32, lora_dropout=0.1)
- test_plugins = [TorchDDPPlugin(), LowLevelZeroPlugin()]
+ test_plugins = [TorchDDPPlugin(), LowLevelZeroPlugin(), HybridParallelPlugin(tp_size=1, pp_size=1)]
test_configs = [
{
"lora_config": lora_config,
@@ -59,6 +60,8 @@ def check_fwd_bwd(model_fn, data_gen_fn, output_transform_fn, loss_fn, task_type
# test fwd bwd correctness
test_model = model_load
+ if isinstance(model_load, HybridParallelModule):
+ model_load = model_load.module.module
model_copy = copy.deepcopy(model_load)
data = data_gen_fn()
diff --git a/tests/test_moe/moe_utils.py b/tests/test_moe/moe_utils.py
index 131932dcb..8c411a33f 100644
--- a/tests/test_moe/moe_utils.py
+++ b/tests/test_moe/moe_utils.py
@@ -1,142 +1,8 @@
import torch
-import torch.distributed as dist
-import torch.nn as nn
-from torch.distributed import ProcessGroup
-from torch.testing import assert_close
-
-from colossalai.booster.plugin.low_level_zero_plugin import LowLevelZeroModel
-from colossalai.legacy.engine.gradient_handler._base_gradient_handler import BaseGradientHandler
-from colossalai.legacy.engine.gradient_handler.utils import bucket_allreduce
-from colossalai.legacy.registry import GRADIENT_HANDLER
-from colossalai.moe.manager import MOE_MANAGER
-from colossalai.moe.utils import get_moe_epsize_param_dict
-
-# from colossalai.shardformer.layer.moe import SparseMLP
-from colossalai.tensor.moe_tensor.api import get_ep_group, get_ep_size, set_moe_tensor_ep_group
-def delete_moe_info(model):
- for _, param in model.named_parameters():
- if hasattr(param, "ep_group"):
- delattr(param, "ep_group")
-
-
-class MoeModel(nn.Module):
- def __init__(self, ep_group: ProcessGroup = None):
- super().__init__()
- self.test_embed = nn.Linear(4, 16, bias=False)
- self.w1 = torch.nn.Parameter(torch.randn(16, 8))
- if ep_group:
- set_moe_tensor_ep_group(self.w1, ep_group)
-
- def forward(self, x):
- x = self.test_embed(x)
- x = torch.matmul(x, self.w1)
-
- return x
-
-
-@GRADIENT_HANDLER.register_module
-class MoeGradientHandler(BaseGradientHandler):
- """A helper class to handle all-reduce operations in a data parallel group and
- moe model parallel. A all-reduce collective communication will be operated in
- :func:`handle_gradient` among a data parallel group.
- For better performance, it bucketizes the gradients of all parameters that are
- the same type to improve the efficiency of communication.
-
- Args:
- model (Module): Model where the gradients accumulate.
- optimizer (Optimizer): Optimizer for updating the parameters.
- """
-
- def __init__(self, model, optimizer=None):
- super().__init__(model, optimizer)
-
- def handle_gradient(self):
- """A method running an all-reduce operation in a data parallel group.
- Then running an all-reduce operation for all parameters in experts
- across moe model parallel group
- """
- if dist.get_world_size() > 1:
- epsize_param_dict = get_moe_epsize_param_dict(self._model)
-
- # epsize is 1, indicating the params are replicated among processes in data parallelism
- # use the ParallelMode.DATA to get data parallel group
- # reduce gradients for all parameters in data parallelism
- if 1 in epsize_param_dict:
- bucket_allreduce(param_list=epsize_param_dict[1])
-
- for ep_size in epsize_param_dict:
- if ep_size != 1 and ep_size != MOE_MANAGER.world_size:
- bucket_allreduce(
- param_list=epsize_param_dict[ep_size], group=MOE_MANAGER.parallel_info_dict[ep_size].dp_group
- )
-
-
-def assert_not_equal_in_group(tensor, process_group=None):
- # all gather tensors from different ranks
- world_size = dist.get_world_size(process_group)
- tensor_list = [torch.empty_like(tensor) for _ in range(world_size)]
- dist.all_gather(tensor_list, tensor, group=process_group)
-
- # check if they are equal one by one
- for i in range(world_size - 1):
- a = tensor_list[i]
- b = tensor_list[i + 1]
- assert not torch.allclose(a, b), (
- f"expected tensors on rank {i} and {i + 1} not to be equal " f"but they are, {a} vs {b}"
- )
-
-
-def run_fwd_bwd(model, data, label, criterion, optimizer, enable_autocast=False):
- model.train()
- with torch.cuda.amp.autocast(enabled=enable_autocast):
- if criterion:
- y = model(data)
- loss = criterion(y, label)
- else:
- loss = model(data, label)
- loss = loss.float()
-
- if isinstance(model, LowLevelZeroModel):
- optimizer.backward(loss)
- else:
- loss.backward()
- return y
-
-
-def sync_local_from_ep(local_model, ep_model, assert_grad_flag: bool = False) -> None:
- """Sync the parameters of tp model from ep model
-
- Args:
- local_model (MoeModule)
- ep_model (MoeModule)
- """
- for (local_name, local_param), (ep_name, ep_param) in zip(
- local_model.named_parameters(), ep_model.named_parameters()
- ):
- if "experts" not in local_name:
- if assert_grad_flag:
- assert torch.allclose(local_param, ep_param), f"local_param: {local_param}, ep_param: {ep_param}"
- assert torch.allclose(local_param.grad, ep_param.grad)
- else:
- local_param.data.copy_(ep_param.data)
- continue
-
- # gather param from ep model
- param_list = [torch.zeros_like(ep_param) for _ in range(get_ep_size(ep_param))]
- dist.all_gather(param_list, ep_param, group=get_ep_group(ep_param))
- all_param = torch.cat(param_list, dim=0)
- if assert_grad_flag:
- grad_list = [torch.zeros_like(ep_param) for _ in range(get_ep_size(ep_param))]
- dist.all_gather(grad_list, ep_param.grad, group=get_ep_group(ep_param))
- all_grad = torch.cat(grad_list, dim=0)
-
- if assert_grad_flag:
- assert torch.allclose(local_param, all_param)
- assert torch.allclose(local_param.grad, all_grad)
- else:
- local_param.data.copy_(all_param.data)
+def assert_loose_close(a, b, dtype: torch.dtype = torch.float32, name=""):
+ assert loose_close(a, b, dtype), f"{name} not close {a.mean()} {b.mean()}"
def loose_close(a, b, dtype: torch.dtype = torch.float32):
@@ -148,8 +14,18 @@ def loose_close(a, b, dtype: torch.dtype = torch.float32):
elif dtype is torch.bfloat16:
rtol = 4e-3
atol = 4e-3
+ else:
+ assert dtype is torch.float32
+ rtol = 1e-05
+ atol = 1e-08
a = a.detach().to(dtype)
b = b.detach().to(dtype).to(a.device)
- assert_close(a, b, rtol=rtol, atol=atol)
+ return torch.allclose(a, b, rtol=rtol, atol=atol)
+
+
+def check_model_equal(model1, model2):
+ assert set(model1.state_dict().keys()) == set(model2.state_dict().keys())
+ for i, ((name, p1), p2) in enumerate(zip(model1.named_parameters(), model2.parameters())):
+ assert_loose_close(p1, p2, p1.dtype)
diff --git a/tests/test_moe/test_deepseek_layer.py b/tests/test_moe/test_deepseek_layer.py
new file mode 100644
index 000000000..d18ba2eac
--- /dev/null
+++ b/tests/test_moe/test_deepseek_layer.py
@@ -0,0 +1,78 @@
+from copy import deepcopy
+
+import pytest
+import torch
+import torch.distributed as dist
+from torch.testing import assert_close
+from transformers import AutoConfig, AutoModel
+
+import colossalai
+from colossalai.booster.plugin.moe_hybrid_parallel_plugin import MoeHybridParallelPlugin
+from colossalai.shardformer.modeling.deepseek import EPDeepseekMoE
+from colossalai.testing.utils import spawn
+
+tokens, n_experts = 7, 4
+hidden_size = 8
+top_k = 2
+
+
+def check_deepseek_moe_layer():
+ torch.cuda.set_device(dist.get_rank())
+ plugin = MoeHybridParallelPlugin(
+ precision="bf16",
+ tp_size=1,
+ pp_size=1,
+ zero_stage=1,
+ ep_size=dist.get_world_size(),
+ )
+
+ config = AutoConfig.from_pretrained(
+ "deepseek-ai/deepseek-moe-16b-base",
+ num_hidden_layers=1,
+ n_routed_experts=n_experts,
+ num_experts_per_tok=top_k,
+ hidden_size=hidden_size,
+ intermediate_size=hidden_size * 2,
+ first_k_dense_replace=0,
+ num_attention_heads=2,
+ trust_remote_code=True,
+ )
+ torch.manual_seed(0)
+ # get the moe layer in auto model
+ orig_model = AutoModel.from_config(config, trust_remote_code=True).layers[0].mlp.cuda()
+ x = torch.rand(1, tokens, hidden_size, requires_grad=True).cuda()
+ orig_output = orig_model(x)
+ model = deepcopy(orig_model)
+ model = EPDeepseekMoE.from_native_module(
+ model,
+ ep_group=plugin.ep_group,
+ moe_dp_group=plugin.moe_dp_group,
+ tp_group=plugin.tp_group,
+ )
+ ep_output = model(x)
+ assert_close(orig_output, ep_output)
+ orig_loss = orig_output.mean()
+ orig_loss.backward()
+ ep_loss = ep_output.mean()
+ ep_loss.backward()
+ assert_close(orig_loss, ep_loss)
+ name_to_p = {n: p for n, p in orig_model.named_parameters()}
+ for n, ep_p in model.named_parameters():
+ p = name_to_p[n]
+ if ep_p.grad is not None:
+ assert_close(p.grad, ep_p.grad)
+
+
+def run_dist(rank: int, world_size: int, port: int):
+ colossalai.launch(rank, world_size, "localhost", port)
+ check_deepseek_moe_layer()
+
+
+@pytest.mark.skip("tested in corresponding sharderformer")
+@pytest.mark.parametrize("world_size", [2])
+def test_deepseek_moe_layer(world_size: int):
+ spawn(run_dist, world_size)
+
+
+if __name__ == "__main__":
+ test_deepseek_moe_layer(2)
diff --git a/tests/test_moe/test_kernel.py b/tests/test_moe/test_kernel.py
index 28e6db441..c81023988 100644
--- a/tests/test_moe/test_kernel.py
+++ b/tests/test_moe/test_kernel.py
@@ -4,8 +4,6 @@ import pytest
import torch
from colossalai.accelerator import get_accelerator
-
-# from colossalai.moe import SparseMLP
from colossalai.moe._operation import MoeCombine, MoeDispatch, moe_cumsum
NUM_EXPERTS = 4
diff --git a/tests/test_moe/test_mixtral_layer.py b/tests/test_moe/test_mixtral_layer.py
index b7b0322e0..bc41ac4f3 100644
--- a/tests/test_moe/test_mixtral_layer.py
+++ b/tests/test_moe/test_mixtral_layer.py
@@ -23,6 +23,7 @@ def check_mixtral_moe_layer():
precision="bf16",
tp_size=1,
pp_size=1,
+ zero_stage=1,
ep_size=dist.get_world_size(),
)
config = MixtralConfig(
@@ -36,7 +37,12 @@ def check_mixtral_moe_layer():
x = torch.rand(1, tokens, hidden_size, requires_grad=True).cuda()
orig_output, orig_logits = orig_model(x)
model = deepcopy(orig_model)
- model = EPMixtralSparseMoeBlock.from_native_module(model, ep_group=plugin.ep_group)
+ model = EPMixtralSparseMoeBlock.from_native_module(
+ model,
+ ep_group=plugin.ep_group,
+ tp_group=plugin.tp_group,
+ moe_dp_group=plugin.moe_dp_group,
+ )
ep_output, ep_logits = model(x)
assert_close(orig_logits, ep_logits)
assert_close(orig_output, ep_output)
@@ -57,7 +63,8 @@ def run_dist(rank: int, world_size: int, port: int):
check_mixtral_moe_layer()
-@pytest.mark.parametrize("world_size", [2, 4])
+@pytest.mark.skip("tested in corresponding sharderformer")
+@pytest.mark.parametrize("world_size", [2])
def test_mixtral_moe_layer(world_size: int):
spawn(run_dist, world_size)
diff --git a/tests/test_moe/test_moe_checkpoint.py b/tests/test_moe/test_moe_checkpoint.py
index 249dd4b97..89f5d1c64 100644
--- a/tests/test_moe/test_moe_checkpoint.py
+++ b/tests/test_moe/test_moe_checkpoint.py
@@ -6,30 +6,23 @@ from copy import deepcopy
import pytest
import torch
import torch.distributed as dist
-from torch.optim import Adam
+from torch.optim import SGD, Adam
from transformers.models.mixtral.configuration_mixtral import MixtralConfig
from transformers.models.mixtral.modeling_mixtral import MixtralForCausalLM
import colossalai
from colossalai.booster import Booster
from colossalai.booster.plugin.moe_hybrid_parallel_plugin import MoeHybridParallelPlugin
-from colossalai.checkpoint_io import MoECheckpointIO
-from colossalai.tensor.moe_tensor.api import is_moe_tensor
+from colossalai.testing import parameterize, spawn
+from colossalai.testing.random import seed_all
from colossalai.testing.utils import spawn
+from tests.test_moe.moe_utils import check_model_equal
tokens, n_experts = 7, 4
hidden_size = 8
top_k = 2
-def check_model_equal(model1, model2):
- assert set(model1.state_dict().keys()) == set(model2.state_dict().keys())
- for i, ((name, p1), p2) in enumerate(zip(model1.named_parameters(), model2.parameters())):
- if not torch.equal(p1.half(), p2.half()):
- print(f"Model parameter {name} is not equal. is_moe_tensor: {is_moe_tensor(p1)}")
- raise AssertionError(f"Model parameter {name} is not equal")
-
-
def get_optimizer_snapshot(optim):
state = {id(k): deepcopy(v) for k, v in optim.state.items()}
param_groups = []
@@ -77,36 +70,44 @@ def check_optimizer_snapshot_equal(snapshot1, snapshot2, param2name, moe_dp_grou
raise AssertionError(f"A total of {count} optim states are not equal")
-def check_mixtral_moe_layer():
+@parameterize(
+ "test_config",
+ [
+ [
+ MixtralConfig(
+ hidden_size=hidden_size,
+ intermediate_size=hidden_size * 2,
+ num_local_experts=n_experts,
+ num_experts_per_tok=top_k,
+ num_attention_heads=2,
+ num_key_value_heads=2,
+ num_hidden_layers=2,
+ ),
+ MixtralForCausalLM,
+ ],
+ ],
+)
+def check_moe_checkpoint(test_config):
+ dtype, precision = torch.float16, "fp16"
+ config, model_cls = test_config
+ torch.cuda.set_device(dist.get_rank())
+
context = tempfile.TemporaryDirectory() if dist.get_rank() == 0 else nullcontext()
with context as f:
- torch.cuda.set_device(dist.get_rank())
if dist.get_rank() == 0:
broadcast_objects = [f] # any picklable object
else:
broadcast_objects = [None]
dist.broadcast_object_list(broadcast_objects, src=0)
- config = MixtralConfig(
- hidden_size=hidden_size,
- intermediate_size=hidden_size * 2,
- num_local_experts=n_experts,
- num_experts_per_tok=top_k,
- num_attention_heads=2,
- num_key_value_heads=2,
- )
- torch.manual_seed(0)
input_ids = torch.randint(0, 100, (2, tokens)).cuda()
- orig_model = MixtralForCausalLM(config).cuda()
+ orig_model = model_cls(config).cuda().to(dtype)
+
+ seed_all(10086)
model = deepcopy(orig_model)
- optimizer = Adam(model.parameters(), lr=1e-3)
+ optimizer = SGD(model.parameters(), lr=1e-3)
plugin = MoeHybridParallelPlugin(
- pp_size=2,
- ep_size=2,
- tp_size=1,
- checkpoint_io=MoECheckpointIO,
- microbatch_size=1,
- zero_stage=1,
+ pp_size=2, ep_size=2, tp_size=1, microbatch_size=1, zero_stage=1, precision=precision
)
booster = Booster(plugin=plugin)
model, optimizer, *_ = booster.boost(model=model, optimizer=optimizer)
@@ -120,7 +121,6 @@ def check_mixtral_moe_layer():
lambda outputs, inputs: outputs.loss,
optimizer,
)
-
tmpdirname = broadcast_objects[0]
model_dir = os.path.join(tmpdirname, "mixtral_model")
hf_model_dir = os.path.join(tmpdirname, "mixtral_hf_model")
@@ -129,13 +129,12 @@ def check_mixtral_moe_layer():
booster.save_model(model, model_dir, shard=True)
dist.barrier()
if dist.get_rank() == 0:
- saved_model = MixtralForCausalLM.from_pretrained(model_dir).cuda()
+ saved_model = model_cls.from_pretrained(model_dir).cuda().to(dtype)
check_model_equal(orig_model, saved_model)
- # check_model_equal(model, saved_model)
saved_model.save_pretrained(hf_model_dir)
dist.barrier()
# check load model
- new_model = MixtralForCausalLM(config).cuda()
+ new_model = model_cls(config).cuda().to(dtype)
new_optimizer = Adam(new_model.parameters(), lr=1e-3)
new_model, new_optimizer, *_ = booster.boost(model=new_model, optimizer=new_optimizer)
booster.load_model(new_model, hf_model_dir)
@@ -163,7 +162,7 @@ def check_mixtral_moe_layer():
def run_dist(rank: int, world_size: int, port: int):
colossalai.launch(rank, world_size, "localhost", port)
- check_mixtral_moe_layer()
+ check_moe_checkpoint()
# Test EP + ZeRO + PP
diff --git a/tests/test_moe/test_moe_ep_tp.py b/tests/test_moe/test_moe_ep_tp.py
index 9bc11033a..e6d2609ee 100644
--- a/tests/test_moe/test_moe_ep_tp.py
+++ b/tests/test_moe/test_moe_ep_tp.py
@@ -1,238 +1,132 @@
-import os
-import warnings
-from typing import Dict
+from copy import deepcopy
import pytest
import torch
import torch.distributed as dist
+from transformers.models.mixtral.configuration_mixtral import MixtralConfig
+from transformers.models.mixtral.modeling_mixtral import MixtralModel
import colossalai
-from colossalai.accelerator import get_accelerator
-from colossalai.moe.manager import MOE_MANAGER
-from colossalai.moe.utils import sync_moe_model_param
+from colossalai.booster.booster import Booster
+from colossalai.booster.plugin import HybridParallelPlugin
+from colossalai.booster.plugin.moe_hybrid_parallel_plugin import MoeHybridParallelPlugin
+from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
+from colossalai.testing.random import seed_all
+from tests.test_moe.moe_utils import assert_loose_close
-# from colossalai.shardformer.layer import SparseMLP
-from colossalai.tensor.moe_tensor.api import get_ep_group, get_ep_rank, get_ep_size, is_moe_tensor
-from colossalai.testing import assert_equal_in_group, rerun_if_address_is_in_use, spawn
-from tests.test_moe.moe_utils import MoeGradientHandler
+NUM_BATCH = 4
+NUM_TOK_PER_BATCH, NUM_EXPERTS = 7, 4
+HIDDEN_SIZE_PER_HEAD = 4
+NUM_HEADS = 4
+TOP_K = 2
-def sync_tp_from_local(tp_model, local_model, assert_grad_flag: bool = False) -> None:
- """Sync the parameters of tp model from local model
+@parameterize("stage", [1])
+@parameterize("ep_size", [2])
+def run_zero_with_original_model(stage: int, ep_size: int):
+ tp_size = dist.get_world_size() // ep_size
+ dtype = torch.bfloat16
- Args:
- tp_model (MoeModule)
- local_model (MoeModule)
- """
- for (tp_name, tp_param), (local_name, local_param) in zip(
- tp_model.named_parameters(), local_model.named_parameters()
- ):
- assert tp_name == local_name
- if not is_moe_tensor(tp_param):
- if assert_grad_flag:
- assert torch.allclose(tp_param, local_param)
- assert torch.allclose(tp_param.grad, local_param.grad)
- else:
- tp_param.data.copy_(local_param.data)
- continue
+ rank = torch.distributed.get_rank()
+ torch.cuda.set_device(dist.get_rank())
- tp_rank = get_ep_rank(tp_param)
- tp_dim = [i for i, (d1, d2) in enumerate(zip(tp_param.shape, local_param.shape)) if d1 != d2][0]
- tp_slice = [slice(None)] * tp_dim + [
- slice(tp_param.shape[tp_dim] * tp_rank, tp_param.shape[tp_dim] * (tp_rank + 1))
- ]
+ seed_all(10086)
- if assert_grad_flag:
- assert torch.allclose(tp_param, local_param[tuple(tp_slice)])
- assert torch.allclose(tp_param.grad, local_param.grad[tuple(tp_slice)])
- else:
- tp_param.data.copy_(local_param[tuple(tp_slice)].data)
-
-
-def sync_tp_from_ep(tp_model, ep_model, assert_grad_flag: bool = False) -> None:
- """Sync the parameters of tp model from ep model
-
- Args:
- tp_model (MoeModule)
- ep_model (MoeModule)
- """
- for (tp_name, tp_param), (ep_name, ep_param) in zip(tp_model.named_parameters(), ep_model.named_parameters()):
- assert tp_name == ep_name
- if not is_moe_tensor(tp_param):
- if assert_grad_flag:
- assert torch.allclose(tp_param, ep_param)
- assert torch.allclose(tp_param.grad, ep_param.grad)
- else:
- tp_param.data.copy_(ep_param.data)
- continue
-
- # gather param from ep model
- param_list = [torch.zeros_like(ep_param) for _ in range(get_ep_size(ep_param))]
- dist.all_gather(param_list, ep_param, group=get_ep_group(ep_param))
- all_param = torch.cat(param_list, dim=0)
- if assert_grad_flag:
- grad_list = [torch.zeros_like(ep_param) for _ in range(get_ep_size(ep_param))]
- dist.all_gather(grad_list, ep_param.grad, group=get_ep_group(ep_param))
- all_grad = torch.cat(grad_list, dim=0)
-
- # get tp param
- tp_dim = [i for i, (d1, d2) in enumerate(zip(tp_param.shape[1:], all_param.shape[1:])) if d1 != d2][0] + 1
- tp_rank = get_ep_rank(tp_param)
- tp_slice = [slice(None)] * tp_dim + [
- slice(tp_param.shape[tp_dim] * tp_rank, tp_param.shape[tp_dim] * (tp_rank + 1))
- ]
- new_tp_param = all_param[tuple(tp_slice)]
- if assert_grad_flag:
- new_grad = all_grad[tuple(tp_slice)]
- if assert_grad_flag:
- assert torch.allclose(tp_param, new_tp_param)
- assert torch.allclose(tp_param.grad, new_grad)
- else:
- tp_param.data.copy_(new_tp_param.data)
-
-
-def sync_local_from_ep(local_model, ep_model, assert_grad_flag: bool = False) -> None:
- """Sync the parameters of tp model from ep model
-
- Args:
- local_model (MoeModule)
- ep_model (MoeModule)
- """
- for (local_name, local_param), (ep_name, ep_param) in zip(
- local_model.named_parameters(), ep_model.named_parameters()
- ):
- assert local_name == ep_name
- if "experts" not in local_name:
- if assert_grad_flag:
- assert torch.allclose(local_param, ep_param)
- assert torch.allclose(local_param.grad, ep_param.grad)
- else:
- local_param.data.copy_(ep_param.data)
- continue
-
- # gather param from ep model
- param_list = [torch.zeros_like(ep_param) for _ in range(get_ep_size(ep_param))]
- dist.all_gather(param_list, ep_param, group=get_ep_group(ep_param))
- all_param = torch.cat(param_list, dim=0)
- if assert_grad_flag:
- grad_list = [torch.zeros_like(ep_param) for _ in range(get_ep_size(ep_param))]
- dist.all_gather(grad_list, ep_param.grad, group=get_ep_group(ep_param))
- all_grad = torch.cat(grad_list, dim=0)
-
- if assert_grad_flag:
- assert torch.allclose(local_param, all_param)
- assert torch.allclose(local_param.grad, all_grad)
- else:
- local_param.data.copy_(all_param.data)
-
-
-def run_test(rank: int, world_size: int, port: int, num_experts: int, batch_size: int, dim: int, config: Dict):
- assert batch_size % world_size == 0
-
- colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
-
- MOE_MANAGER.__init__()
- MOE_MANAGER.setup(parallel=None)
- local_model = SparseMLP(num_experts=num_experts, hidden_size=dim, intermediate_size=dim * 2)
- MOE_MANAGER.__init__()
- MOE_MANAGER.setup(parallel="EP")
- enable_hierarchical_comm = config.get("enable_hierarchical_comm", False)
- if enable_hierarchical_comm:
- os.environ["LOCAL_WORLD_SIZE"] = str(world_size)
- ep_model = SparseMLP(
- num_experts=num_experts,
- hidden_size=dim,
- intermediate_size=dim * 2,
- enable_hierarchical_comm=enable_hierarchical_comm,
+ config = MixtralConfig(
+ hidden_size=HIDDEN_SIZE_PER_HEAD * NUM_HEADS,
+ intermediate_size=HIDDEN_SIZE_PER_HEAD * NUM_HEADS * 2,
+ num_hidden_layers=2,
+ num_attention_heads=NUM_HEADS,
+ num_key_value_heads=NUM_HEADS,
+ num_local_experts=NUM_EXPERTS,
+ num_experts_per_tok=TOP_K,
)
- MOE_MANAGER.__init__()
- MOE_MANAGER.setup(parallel="TP")
- tp_model = SparseMLP(num_experts=num_experts, hidden_size=dim, intermediate_size=dim * 2)
- ep_model = ep_model.to(get_accelerator().get_current_device())
- tp_model = tp_model.to(get_accelerator().get_current_device())
- local_model = local_model.to(get_accelerator().get_current_device())
+ torch_model = MixtralModel(config).to(dtype).cuda()
- # sync ep param
- sync_moe_model_param(ep_model)
- dist_dict = MOE_MANAGER.parallel_info_dict
- assert_equal_in_group(ep_model.experts.wi.data, dist_dict[world_size].dp_group)
- assert_equal_in_group(ep_model.experts.wo.data, dist_dict[world_size].dp_group)
- ep_grad_handler = MoeGradientHandler(ep_model)
- # sync local param
- sync_local_from_ep(local_model, ep_model)
- # sync tp param
- sync_tp_from_ep(tp_model, ep_model)
- tp_grad_handler = MoeGradientHandler(tp_model)
-
- rank = dist.get_rank()
- input_data = torch.randn(batch_size, dim, device=get_accelerator().get_current_device())
- micro_batch_size = batch_size // world_size
- index = rank * micro_batch_size
- # NOTE: ep & tp takes in sharded data for each process
- shard_data = input_data.detach()[index : index + micro_batch_size]
-
- out_local = local_model(input_data)
- MOE_MANAGER.reset_loss()
- out_tp = tp_model(shard_data)
- MOE_MANAGER.reset_loss()
- out_ep = ep_model(shard_data)
- MOE_MANAGER.reset_loss()
-
- assert torch.allclose(
- out_tp, out_ep, atol=1e-6
- ), f"Rank {rank} failed, max diff: {torch.max(torch.abs(out_tp - out_ep))}"
- try:
- out_local_slice = out_local[index : index + micro_batch_size]
- assert torch.allclose(
- out_ep, out_local_slice, atol=1e-6
- ), f"Rank {rank} failed, max diff: {torch.max(torch.abs(out_ep - out_local_slice))}"
- except AssertionError:
- """
- e.g., in local model, tokens = 4, capacity = 2, experts = 2, topk = 1
- router yields [01] --> [0], [23] --> [1], this is valid as capacity is 2
- However, in ep mode, there are 2 separate routers dealing with sharded data.
- Assume router 0 handles token [01] and router 1 handles token [23].
- Note that for each router the capacity is only 1 !!!
- Thus, router 0 may yields [0] --> [0] or [1] --> [0], but not both.
- The same thing happens on router 1. And finally some tokens are dropped due to the sharded nature.
- """
- warnings.warn(
- "EP & TP may result in different behavior from local model. " "Please check the comments for details."
+ zero_model = deepcopy(torch_model).to(dtype)
+ zero_optimizer = torch.optim.SGD(zero_model.parameters(), lr=1)
+ moe_booster = Booster(
+ plugin=MoeHybridParallelPlugin(
+ tp_size=tp_size,
+ moe_tp_size=tp_size,
+ pp_size=1,
+ ep_size=ep_size,
+ zero_stage=stage,
+ overlap_communication=False,
+ initial_scale=1,
)
+ )
+ zero_model, zero_optimizer, _, _, _ = moe_booster.boost(zero_model, zero_optimizer)
- out_local.mean().backward()
- out_tp.mean().backward()
- tp_grad_handler.handle_gradient()
- out_ep.mean().backward()
- ep_grad_handler.handle_gradient()
-
- assert_equal_in_group(ep_model.experts.wi.grad, dist_dict[world_size].dp_group)
- assert_equal_in_group(ep_model.experts.wo.grad, dist_dict[world_size].dp_group)
- sync_tp_from_ep(tp_model, ep_model, assert_grad_flag=True)
- try:
- sync_local_from_ep(local_model, ep_model, assert_grad_flag=True)
- except AssertionError:
- warnings.warn(
- "EP & TP may result in different behavior from local model. " "Please check the comments for details."
+ hybird_booster = Booster(
+ plugin=HybridParallelPlugin(
+ tp_size=tp_size,
+ pp_size=1,
+ zero_stage=stage,
+ overlap_communication=False,
+ initial_scale=1,
)
+ )
+ hybrid_model, hybrid_optimizer, _, _, _ = hybird_booster.boost(
+ torch_model, torch.optim.SGD(torch_model.parameters(), lr=1)
+ )
+ # create different input
+ seed_all(1453 + rank)
+
+ hybrid_model.train()
+ zero_model.train()
+ for _ in range(2):
+ # zero-dp forward
+ input_data = torch.rand(
+ NUM_BATCH, NUM_TOK_PER_BATCH, HIDDEN_SIZE_PER_HEAD * NUM_HEADS, requires_grad=True
+ ).cuda()
+ zero_output = zero_model(inputs_embeds=input_data.to(dtype)).last_hidden_state.mean()
+ # zero-dp backward
+ zero_optimizer.backward(zero_output)
+ # torch-ddp forward
+ hybrid_output = hybrid_model(inputs_embeds=input_data.to(dtype)).last_hidden_state.mean()
+ assert_loose_close(zero_output, hybrid_output, dtype=dtype)
+ # torch-ddp backward
+ hybrid_optimizer.backward(hybrid_output)
+
+ # check grad
+ name_to_p = {n: p for n, p in hybrid_model.named_parameters()}
+ for n, p in zero_model.named_parameters():
+ zero_grad = zero_optimizer.get_param_grad(p)
+ if name_to_p[n].grad is None:
+ name_to_p[n].grad = torch.zeros_like(name_to_p[n])
+ continue
+ if zero_grad.shape != name_to_p[n].grad.shape: # TODO check sharded and sliced moe
+ continue
+ assert_loose_close(zero_grad, name_to_p[n].grad, dtype=dtype, name=n)
+
+ # zero-dp step
+ zero_optimizer.step()
+
+ # original model step
+ hybrid_optimizer.step()
+
+ # check updated param
+ for n, p in zero_model.named_parameters():
+ if p.data.shape != name_to_p[n].data.shape: # TODO check sharded and sliced moe
+ continue
+ assert_loose_close(p.data, name_to_p[n].data, dtype=dtype, name=n)
+
+ print(f"{dist.get_rank()} test passed")
-@pytest.mark.skip(reason="moe need to be refactored")
+def run_dist(rank, world_size, port):
+ colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
+ run_zero_with_original_model()
+
+
+@pytest.mark.skip("tested in corresponding sharderformer")
@pytest.mark.dist
-@pytest.mark.parametrize("num_experts", [4, 64])
-@pytest.mark.parametrize("batch_size", [16])
-@pytest.mark.parametrize("dim", [64])
-@pytest.mark.parametrize(
- "config",
- [
- {"enable_hierarchical_comm": False},
- {"enable_hierarchical_comm": True},
- ],
-)
+@pytest.mark.parametrize("world_size", [4])
@rerun_if_address_is_in_use()
-def test_moe_ep_tp(num_experts: int, batch_size: int, dim: int, config: Dict):
- spawn(run_test, 2, num_experts=num_experts, batch_size=batch_size, dim=dim, config=config)
+def test_moe_ep_tp(world_size):
+ spawn(run_dist, world_size)
if __name__ == "__main__":
- test_moe_ep_tp(num_experts=8, batch_size=32, dim=32)
+ test_moe_ep_tp(world_size=4)
diff --git a/tests/test_moe/test_moe_ep_zero.py b/tests/test_moe/test_moe_ep_zero.py
new file mode 100644
index 000000000..2d4e638b6
--- /dev/null
+++ b/tests/test_moe/test_moe_ep_zero.py
@@ -0,0 +1,119 @@
+from copy import deepcopy
+
+import pytest
+import torch
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+from transformers.models.mixtral.configuration_mixtral import MixtralConfig
+from transformers.models.mixtral.modeling_mixtral import MixtralModel
+
+import colossalai
+from colossalai.booster.booster import Booster
+from colossalai.booster.plugin.moe_hybrid_parallel_plugin import MoeHybridParallelPlugin
+from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
+from colossalai.testing.random import seed_all
+from tests.test_moe.moe_utils import assert_loose_close
+
+NUM_BATCH = 4
+NUM_TOK_PER_BATCH, NUM_EXPERTS = 7, 4
+HIDDEN_SIZE_PER_HEAD = 4
+NUM_HEADS = 2
+TOP_K = 1
+
+
+@parameterize("stage", [1])
+@parameterize("ep_size", [2, 4])
+def run_zero_with_original_model(stage: int, ep_size: int):
+ dtype = torch.bfloat16
+
+ rank = torch.distributed.get_rank()
+ torch.cuda.set_device(dist.get_rank())
+
+ plugin = MoeHybridParallelPlugin(
+ pp_size=1, tp_size=1, ep_size=ep_size, zero_stage=stage, overlap_communication=False, initial_scale=1
+ )
+ booster = Booster(plugin=plugin)
+
+ seed_all(10086)
+
+ config = MixtralConfig(
+ hidden_size=HIDDEN_SIZE_PER_HEAD * NUM_HEADS,
+ intermediate_size=HIDDEN_SIZE_PER_HEAD * NUM_HEADS * 2,
+ num_hidden_layers=2,
+ num_attention_heads=NUM_HEADS,
+ num_key_value_heads=NUM_HEADS,
+ num_local_experts=NUM_EXPERTS,
+ num_experts_per_tok=TOP_K,
+ )
+
+ torch_model = MixtralModel(config).to(dtype).cuda()
+
+ zero_model = deepcopy(torch_model).to(dtype)
+ zero_optimizer = torch.optim.SGD(zero_model.parameters(), lr=1)
+
+ zero_model, zero_optimizer, _, _, _ = booster.boost(zero_model, zero_optimizer)
+
+ ddp_model = DDP(
+ torch_model.cuda(),
+ process_group=plugin.dp_group,
+ find_unused_parameters=True, # important for torch ddp, not all experts are routed
+ ).cuda()
+ ddp_optimizer = torch.optim.SGD(ddp_model.parameters(), lr=1)
+
+ # create different input
+ seed_all(1453 + rank)
+
+ ddp_model.train()
+ zero_model.train()
+ for _ in range(2):
+ # zero-dp forward
+ input_data = torch.rand(
+ NUM_BATCH, NUM_TOK_PER_BATCH, HIDDEN_SIZE_PER_HEAD * NUM_HEADS, requires_grad=True
+ ).cuda()
+ zero_output = zero_model(inputs_embeds=input_data.to(dtype)).last_hidden_state.mean()
+ # zero-dp backward
+ zero_optimizer.backward(zero_output)
+
+ # torch-ddp forward
+ ddp_output = ddp_model(inputs_embeds=input_data.to(dtype)).last_hidden_state.mean()
+ assert_loose_close(zero_output, ddp_output, dtype=dtype)
+ # torch-ddp backward
+ ddp_output.backward()
+
+ # check grad
+ name_to_p = {n: p for n, p in ddp_model.named_parameters()}
+ for n, p in zero_model.named_parameters():
+ zero_grad = zero_optimizer.get_param_grad(p)
+ if name_to_p[n].grad is None:
+ name_to_p[n].grad = torch.zeros_like(name_to_p[n].data)
+ continue
+ assert_loose_close(zero_grad, name_to_p[n].grad, dtype=dtype, name=n)
+
+ # zero-dp step
+ zero_optimizer.step()
+
+ # original model step
+ ddp_optimizer.step()
+
+ # check updated param
+ for n, p in zero_model.named_parameters():
+ assert_loose_close(p.data, name_to_p[n].data, dtype=dtype, name=n)
+
+ print(f"{dist.get_rank()} test passed")
+
+
+def run_dist(rank, world_size, port):
+ colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
+ run_zero_with_original_model()
+
+
+@pytest.mark.skip("tested in corresponding sharderformer")
+@pytest.mark.dist
+@pytest.mark.parametrize("world_size", [4])
+@rerun_if_address_is_in_use()
+def test_moe_ep_zero(world_size):
+ spawn(run_dist, world_size)
+
+
+if __name__ == "__main__":
+ test_moe_ep_zero(world_size=4)
diff --git a/tests/test_moe/test_moe_zero_fwd_bwd_optim.py b/tests/test_moe/test_moe_zero_fwd_bwd_optim.py
deleted file mode 100644
index 042b3d8ae..000000000
--- a/tests/test_moe/test_moe_zero_fwd_bwd_optim.py
+++ /dev/null
@@ -1,132 +0,0 @@
-from copy import deepcopy
-
-import pytest
-import torch
-import torch.distributed as dist
-from torch.nn.parallel import DistributedDataParallel as DDP
-from transformers.models.mixtral.configuration_mixtral import MixtralConfig
-from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock
-
-import colossalai
-from colossalai.booster.plugin.moe_hybrid_parallel_plugin import MoeHybridParallelPlugin
-from colossalai.shardformer.modeling.mixtral import EPMixtralSparseMoeBlock
-from colossalai.tensor.moe_tensor.api import is_moe_tensor
-from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
-from colossalai.testing.random import seed_all
-from colossalai.zero import LowLevelZeroOptimizer
-from tests.test_moe.moe_utils import loose_close
-
-tokens, n_experts = 7, 4
-hidden_size = 8
-top_k = 2
-
-
-def split_grad(grad, world_size):
- with torch.no_grad():
- grad = grad.clone().detach().flatten()
- padding_size = (world_size - grad.numel() % world_size) % world_size
- if padding_size > 0:
- grad = torch.nn.functional.pad(grad, [0, padding_size])
- splited_grad = grad.split(grad.numel() // world_size)
- return splited_grad
-
-
-@parameterize("dtype", [torch.float16, torch.bfloat16])
-@parameterize("master_weights", [True, False])
-@parameterize("stage", [1, 2])
-def run_zero_with_original_model(world_size, master_weights: bool, dtype: torch.dtype, stage: int):
- rank = torch.distributed.get_rank()
- torch.cuda.set_device(dist.get_rank())
- plugin = MoeHybridParallelPlugin(
- tp_size=1,
- pp_size=1,
- ep_size=dist.get_world_size() // 2,
- )
-
- seed_all(10086)
- config = MixtralConfig(
- hidden_size=hidden_size,
- intermediate_size=hidden_size * 2,
- num_local_experts=n_experts,
- num_experts_per_tok=top_k,
- )
-
- orig_model = MixtralSparseMoeBlock(config).to(dtype).cuda()
-
- ori_model = DDP(orig_model.cuda(), static_graph=True).cuda()
-
- zero_model = deepcopy(orig_model).to(dtype)
- zero_model = EPMixtralSparseMoeBlock.from_native_module(zero_model, ep_group=plugin.ep_group)
-
- zero_optimizer = torch.optim.SGD(zero_model.parameters(), lr=1)
- pg_param_list = {plugin.global_dp_group: [], plugin.moe_dp_group: []}
- for p in zero_model.parameters():
- if is_moe_tensor(p):
- pg_param_list[plugin.moe_dp_group].append(p)
- else:
- pg_param_list[plugin.global_dp_group].append(p)
-
- zero_optimizer = LowLevelZeroOptimizer(
- zero_optimizer,
- pg_to_param_list=pg_param_list,
- master_weights=master_weights,
- initial_scale=1,
- overlap_communication=False,
- partition_grad=True,
- )
-
- ori_optimizer = torch.optim.SGD(ori_model.parameters(), lr=1)
-
- # create
- seed_all(1453 + rank)
-
- for _ in range(2):
- # zero-dp forward
- input_data = torch.rand(1, tokens, hidden_size).cuda()
- zero_output, zero_logits = zero_model(input_data.to(dtype))
-
- # torch-ddp forward
- ori_output, ori_logits = ori_model(input_data.to(dtype))
- loose_close(zero_output, ori_output, dtype=dtype)
-
- # zero-dp backward
- zero_optimizer.backward(zero_output.mean().float())
-
- # torch-ddp backward
- ori_output.mean().backward()
-
- # check grad
- name_to_p = {n: p for n, p in ori_model.module.named_parameters()}
- for n, p in zero_model.named_parameters():
- zero_grad = zero_optimizer.get_param_grad(p)
- if name_to_p[n].grad is None:
- assert zero_grad is None
- continue
-
- loose_close(zero_grad, name_to_p[n].grad, dtype=dtype)
-
- # zero-dp step
- zero_optimizer.step()
-
- # original model step
- ori_optimizer.step()
-
- # check updated param
- for n, p in zero_model.named_parameters():
- loose_close(p.data, name_to_p[n].data, dtype=dtype)
-
-
-def run_dist(rank, world_size, port):
- colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
- run_zero_with_original_model(world_size=world_size)
-
-
-@pytest.mark.dist
-@pytest.mark.parametrize("world_size", [2, 4])
-@rerun_if_address_is_in_use()
-def test_moe_zero_model(world_size):
- spawn(run_dist, world_size)
-
-
-if __name__ == "__main__":
- test_moe_zero_model(world_size=4)
diff --git a/tests/test_shardformer/test_model/_utils.py b/tests/test_shardformer/test_model/_utils.py
index 1ffcc541a..190fee129 100644
--- a/tests/test_shardformer/test_model/_utils.py
+++ b/tests/test_shardformer/test_model/_utils.py
@@ -1,6 +1,6 @@
import copy
from contextlib import nullcontext
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional, Type
import torch
import torch.distributed as dist
@@ -117,7 +117,12 @@ def check_state_dict(org_model: Module, sharded_model: Module, name: str = ""):
def build_model_from_hybrid_plugin(
- model_fn: Callable, loss_fn: Callable, test_config: Dict[str, Any], optim_class=Adam, sharded_optim_class=Adam
+ model_fn: Callable,
+ loss_fn: Callable,
+ test_config: Dict[str, Any],
+ optim_class=Adam,
+ sharded_optim_class=Adam,
+ pluggin_cls: Type[HybridParallelPlugin] = HybridParallelPlugin,
):
use_lazy_init = False
if "use_lazy_init" in test_config:
@@ -149,9 +154,10 @@ def build_model_from_hybrid_plugin(
else:
org_optimizer = optim_class(org_model.parameters(), lr=1e-3)
sharded_optimizer = sharded_optim_class(sharded_model.parameters(), lr=1e-3)
+
criterion = loss_fn
- plugin = HybridParallelPlugin(**test_config)
+ plugin = pluggin_cls(**test_config)
booster = Booster(plugin=plugin)
sharded_model, sharded_optimizer, criterion, _, _ = booster.boost(sharded_model, sharded_optimizer, criterion)
diff --git a/tests/test_shardformer/test_model/test_shard_chatglm2.py b/tests/test_shardformer/test_model/test_shard_chatglm2.py
index 6ce020b68..92c077950 100644
--- a/tests/test_shardformer/test_model/test_shard_chatglm2.py
+++ b/tests/test_shardformer/test_model/test_shard_chatglm2.py
@@ -136,6 +136,44 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,
@parameterize(
"test_config",
[
+ { # Ulysess + Flash attention
+ "tp_size": 1,
+ "pp_size": 2,
+ "sp_size": 2,
+ "num_microbatches": 2,
+ "enable_sequence_parallelism": True,
+ "sequence_parallelism_mode": "all_to_all",
+ "enable_flash_attention": True,
+ "use_lazy_init": True,
+ "zero_stage": 1,
+ "precision": "fp16",
+ "initial_scale": 1,
+ },
+ {
+ "tp_size": 2,
+ "pp_size": 2,
+ "sp_size": 2,
+ "num_microbatches": 2,
+ "enable_sequence_parallelism": True,
+ "sequence_parallelism_mode": "split_gather",
+ "enable_flash_attention": True,
+ "use_lazy_init": True,
+ "zero_stage": 1,
+ "precision": "fp16",
+ "initial_scale": 1,
+ },
+ {
+ "tp_size": 1,
+ "pp_size": 1,
+ "sp_size": 2,
+ "num_microbatches": 1,
+ "enable_sequence_parallelism": True,
+ "sequence_parallelism_mode": "all_to_all",
+ "use_lazy_init": True,
+ "zero_stage": 1,
+ "precision": "fp16",
+ "initial_scale": 1,
+ },
{
"tp_size": 4,
"pp_size": 1,
diff --git a/tests/test_shardformer/test_model/test_shard_command.py b/tests/test_shardformer/test_model/test_shard_command.py
index 4d66692a4..3281b50e1 100644
--- a/tests/test_shardformer/test_model/test_shard_command.py
+++ b/tests/test_shardformer/test_model/test_shard_command.py
@@ -58,6 +58,7 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,
# Check the grad when using ZeRO-1 and ZeRO-2
if (
booster.plugin.zero_stage in [1, 2]
+ and booster.plugin.shard_config.pipeline_stage_manager is None
and booster.plugin.shard_config.enable_sequence_parallelism
and booster.plugin.shard_config.sequence_parallelism_mode == "all_to_all"
):
@@ -154,6 +155,45 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,
@parameterize(
"test_config",
[
+ { # Ulysess + Flash attention
+ "tp_size": 1,
+ "pp_size": 2,
+ "sp_size": 2,
+ "num_microbatches": 2,
+ "enable_sequence_parallelism": True,
+ "sequence_parallelism_mode": "all_to_all",
+ "enable_flash_attention": True,
+ "use_lazy_init": True,
+ "zero_stage": 1,
+ "precision": "fp16",
+ "initial_scale": 1,
+ },
+ {
+ "tp_size": 2,
+ "pp_size": 2,
+ "sp_size": 2,
+ "num_microbatches": 2,
+ "enable_sequence_parallelism": True,
+ "sequence_parallelism_mode": "split_gather",
+ "enable_flash_attention": True,
+ "use_lazy_init": True,
+ "zero_stage": 1,
+ "precision": "fp16",
+ "initial_scale": 1,
+ },
+ {
+ "tp_size": 2,
+ "pp_size": 2,
+ "sp_size": 2,
+ "num_microbatches": 2,
+ "enable_sequence_parallelism": True,
+ "sequence_parallelism_mode": "ring",
+ "enable_flash_attention": True,
+ "use_lazy_init": True,
+ "zero_stage": 1,
+ "precision": "fp16",
+ "initial_scale": 1,
+ },
{
"tp_size": 2,
"pp_size": 1,
diff --git a/tests/test_shardformer/test_model/test_shard_deepseek.py b/tests/test_shardformer/test_model/test_shard_deepseek.py
new file mode 100644
index 000000000..46da4522f
--- /dev/null
+++ b/tests/test_shardformer/test_model/test_shard_deepseek.py
@@ -0,0 +1,196 @@
+import os
+import shutil
+from copy import deepcopy
+from typing import Tuple
+
+import pytest
+import torch
+import torch.distributed
+import torch.distributed as dist
+from transformers import AutoConfig, AutoModel
+
+import colossalai
+from colossalai.booster.booster import Booster
+from colossalai.booster.plugin.moe_hybrid_parallel_plugin import MoeHybridParallelPlugin
+from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
+from colossalai.testing.random import seed_all
+from tests.test_moe.moe_utils import assert_loose_close, check_model_equal
+
+NUM_BATCH = 8
+NUM_TOK_PER_BATCH, NUM_EXPERTS = 4, 2
+NUM_LAYERS = 4
+HIDDEN_SIZE_PER_HEAD = 4
+NUM_HEADS = 4
+TOP_K = 2
+
+
+CHECKED_CONFIG = [ # FOR_WORLD=4
+ (1, 4, 1, 1, 1),
+ (1, 1, 4, 1, 1),
+ (1, 1, 1, 4, 1),
+ (1, 1, 1, 1, 4),
+ (0, 1, 4, 1, 1),
+ (0, 1, 1, 4, 1),
+ (0, 1, 1, 1, 4),
+ (1, 2, 1, 1, 1),
+]
+
+
+@parameterize(
+ "config",
+ [
+ (1, 2, 2, 1, 1),
+ (1, 2, 1, 2, 1),
+ (1, 2, 1, 1, 2),
+ ],
+)
+def run_zero_with_original_model(config: Tuple[int, ...]):
+ stage, ep_size, pp_size, tp_size, sp_size = config
+ world_size = dist.get_world_size()
+ rank = dist.get_rank()
+ dtype, precision = torch.float16, "fp16"
+ torch.cuda.set_device(dist.get_rank())
+
+ plugin = MoeHybridParallelPlugin(
+ pp_size=pp_size,
+ num_microbatches=pp_size,
+ tp_size=tp_size,
+ sp_size=sp_size,
+ ep_size=ep_size,
+ zero_stage=stage,
+ enable_sequence_parallelism=sp_size > 1,
+ sequence_parallelism_mode="all_to_all" if sp_size > 1 else None,
+ enable_flash_attention=sp_size > 1,
+ overlap_communication=False,
+ initial_scale=1,
+ precision=precision,
+ find_unused_parameters=True,
+ )
+ dp_size = plugin.dp_size
+
+ booster = Booster(plugin=plugin)
+
+ assert pp_size <= NUM_LAYERS, "pp_size should be less than or equal to NUM_LAYERS"
+ config = AutoConfig.from_pretrained(
+ "deepseek-ai/deepseek-moe-16b-base",
+ hidden_size=HIDDEN_SIZE_PER_HEAD * NUM_HEADS,
+ intermediate_size=HIDDEN_SIZE_PER_HEAD * NUM_HEADS * 2,
+ moe_intermediate_size=HIDDEN_SIZE_PER_HEAD * NUM_HEADS * 2,
+ num_hidden_layers=4,
+ num_attention_heads=NUM_HEADS,
+ num_key_value_heads=NUM_HEADS,
+ first_k_dense_replace=1,
+ attn_implementation="flash_attention_2",
+ torch_dtype="float16",
+ n_routed_experts=NUM_EXPERTS,
+ num_experts_per_tok=TOP_K,
+ trust_remote_code=True,
+ )
+
+ # init model with the same seed
+ seed_all(10086)
+
+ torch_model = AutoModel.from_config(config, trust_remote_code=True).cuda().to(dtype)
+ torch_optimizer = torch.optim.SGD(torch_model.parameters(), lr=1)
+
+ parallel_model = deepcopy(torch_model)
+ parallel_optimizer = torch.optim.SGD(parallel_model.parameters(), lr=1)
+ parallel_model, parallel_optimizer, _, _, _ = booster.boost(parallel_model, parallel_optimizer)
+
+ # create different input along dp axis
+ seed_all(1453 + rank)
+
+ torch_model.train()
+ parallel_model.train()
+ for _ in range(2):
+ # gen random input
+ input_embeddings = torch.rand(
+ NUM_BATCH, NUM_TOK_PER_BATCH, HIDDEN_SIZE_PER_HEAD * NUM_HEADS, requires_grad=True
+ ).cuda()
+ dist.all_reduce(
+ input_embeddings, group=plugin.pp_group
+ ) # pp inputs except the first stage doesn't matter, but need to be replicate for torch model check
+
+ dist.all_reduce(input_embeddings, group=plugin.tp_group) # tp group duplicate input
+ dist.all_reduce(input_embeddings, group=plugin.sp_group) # sp group duplicate input
+
+ # run the model with hybrid parallel
+ if booster.plugin.stage_manager is not None:
+ # for test with pp
+ data_iter = iter([{"inputs_embeds": input_embeddings}])
+ sharded_output = booster.execute_pipeline(
+ data_iter,
+ parallel_model,
+ lambda x, y: x[0].mean(),
+ parallel_optimizer,
+ return_loss=True,
+ return_outputs=True,
+ )
+ if booster.plugin.stage_manager.is_last_stage():
+ parallel_output = sharded_output["loss"]
+ else:
+ parallel_output = torch.tensor(12345.0, device="cuda")
+
+ # broadcast along pp axis
+ dist.broadcast(
+ parallel_output, src=dist.get_process_group_ranks(plugin.pp_group)[-1], group=plugin.pp_group
+ )
+ else:
+ # for test without pp
+ parallel_output = parallel_model(inputs_embeds=input_embeddings.to(dtype)).last_hidden_state.mean()
+ parallel_optimizer.backward(parallel_output)
+ parallel_optimizer.step()
+ parallel_optimizer.zero_grad()
+ dist.all_reduce(parallel_output, group=plugin.dp_group)
+
+ # ===================================================================================
+ # run normal model with all dp(different) inputs
+ all_inputs = [torch.empty_like(input_embeddings) for _ in range(dp_size)]
+ dist.all_gather(all_inputs, input_embeddings, group=plugin.dp_group)
+ torch_output_sum = 0
+ for input_data_ in all_inputs:
+ torch_output = torch_model(inputs_embeds=input_data_.to(dtype)).last_hidden_state.mean()
+ torch_output.backward()
+ torch_output_sum += torch_output.detach()
+ # avg dp grads follows zero optimizer
+ for p in torch_model.parameters():
+ if p.grad is not None:
+ p.grad /= dp_size
+ torch_optimizer.step()
+ torch_optimizer.zero_grad()
+
+ assert_loose_close(parallel_output, torch_output_sum, dtype=dtype)
+
+ # use checkpoint to load sharded zero model
+ model_dir = "./test_deepseek"
+ if rank == world_size - 1:
+ os.makedirs(model_dir, exist_ok=True)
+
+ dist.barrier()
+ booster.save_model(parallel_model, model_dir, shard=True)
+ dist.barrier()
+
+ saved_model = AutoModel.from_pretrained(model_dir, trust_remote_code=True).cuda()
+ check_model_equal(torch_model, saved_model)
+ dist.barrier()
+
+ if rank == world_size - 1:
+ shutil.rmtree(model_dir)
+
+ print(f"rank {dist.get_rank()} test passed")
+
+
+def run_dist(rank, world_size, port):
+ colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
+ run_zero_with_original_model()
+
+
+@pytest.mark.dist
+@pytest.mark.parametrize("world_size", [4])
+@rerun_if_address_is_in_use()
+def test_deepseek(world_size):
+ spawn(run_dist, world_size)
+
+
+if __name__ == "__main__":
+ test_deepseek(world_size=4)
diff --git a/tests/test_shardformer/test_model/test_shard_llama.py b/tests/test_shardformer/test_model/test_shard_llama.py
index 8fe18f69b..88e54176b 100644
--- a/tests/test_shardformer/test_model/test_shard_llama.py
+++ b/tests/test_shardformer/test_model/test_shard_llama.py
@@ -59,10 +59,12 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,
if (
booster.plugin.zero_stage in [1, 2]
and booster.plugin.shard_config.enable_sequence_parallelism
+ and booster.plugin.shard_config.pipeline_stage_manager is None
and booster.plugin.shard_config.sequence_parallelism_mode == "all_to_all"
):
+ master2working = sharded_optimizer.get_master_to_working_map()
for p1, p2 in zip(llama_model.parameters(), sharded_optimizer._master_param_groups_of_current_rank[0]):
- working_p = sharded_optimizer.master_to_working_param[id(p2)]
+ working_p = master2working[id(p2)]
grads = sharded_optimizer.get_partitioned_gradients_by_param_id(0, id(working_p))
grad_index = (
0
@@ -146,6 +148,19 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,
@parameterize(
"test_config",
[
+ { # Ulysess + Flash attention
+ "tp_size": 1,
+ "pp_size": 2,
+ "sp_size": 2,
+ "num_microbatches": 2,
+ "enable_sequence_parallelism": True,
+ "sequence_parallelism_mode": "all_to_all",
+ "enable_flash_attention": True,
+ "use_lazy_init": True,
+ "zero_stage": 0,
+ "precision": "fp16",
+ "initial_scale": 1,
+ },
{ # Test ring + Flash attention
"tp_size": 2,
"pp_size": 1,
@@ -159,19 +174,6 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,
"precision": "fp16",
"initial_scale": 1,
},
- { # Ulysess + Flash attention
- "tp_size": 1,
- "pp_size": 2,
- "sp_size": 2,
- "num_microbatches": 2,
- "enable_sequence_parallelism": True,
- "sequence_parallelism_mode": "all_to_all",
- "enable_flash_attention": True,
- "use_lazy_init": True,
- "zero_stage": 1,
- "precision": "fp16",
- "initial_scale": 1,
- },
{
"tp_size": 1,
"pp_size": 1,
@@ -245,7 +247,6 @@ def run_llama_test(test_config):
except Exception as e:
print(f"Failed config: {test_config}")
raise e
-
clear_layout_converter()
Randomizer.reset_index()
torch.cuda.empty_cache()
diff --git a/tests/test_shardformer/test_model/test_shard_mixtral.py b/tests/test_shardformer/test_model/test_shard_mixtral.py
new file mode 100644
index 000000000..de09eedcb
--- /dev/null
+++ b/tests/test_shardformer/test_model/test_shard_mixtral.py
@@ -0,0 +1,190 @@
+import os
+import shutil
+from copy import deepcopy
+from typing import Tuple
+
+import pytest
+import torch
+import torch.distributed
+import torch.distributed as dist
+from transformers.models.mixtral.configuration_mixtral import MixtralConfig
+from transformers.models.mixtral.modeling_mixtral import MixtralModel
+
+import colossalai
+from colossalai.booster.booster import Booster
+from colossalai.booster.plugin.moe_hybrid_parallel_plugin import MoeHybridParallelPlugin
+from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
+from colossalai.testing.random import seed_all
+from tests.test_moe.moe_utils import assert_loose_close, check_model_equal
+
+NUM_BATCH = 8
+NUM_TOK_PER_BATCH, NUM_EXPERTS = 4, 4
+NUM_LAYERS = 4
+HIDDEN_SIZE_PER_HEAD = 4
+NUM_HEADS = 4
+TOP_K = 1
+
+CHECKED_CONFIG = [ # FOR WORLD=4
+ (0, 1, 4, 1, 1),
+ (0, 1, 1, 4, 1),
+ (0, 1, 1, 1, 4),
+ (1, 4, 1, 1, 1),
+ (1, 1, 4, 1, 1),
+ (1, 1, 1, 4, 1),
+ (1, 1, 1, 1, 4),
+ (1, 2, 1, 1, 1),
+]
+
+
+@parameterize(
+ "config",
+ [
+ (1, 2, 2, 1, 1),
+ (1, 2, 1, 2, 1),
+ (1, 2, 1, 1, 2),
+ ],
+)
+def run_zero_with_original_model(config: Tuple[int, ...]):
+ stage, ep_size, pp_size, tp_size, sp_size = config
+ world_size = dist.get_world_size()
+ rank = dist.get_rank()
+ dtype, precision = torch.float16, "fp16"
+ torch.cuda.set_device(dist.get_rank())
+
+ plugin = MoeHybridParallelPlugin(
+ pp_size=pp_size,
+ num_microbatches=pp_size,
+ tp_size=tp_size,
+ sp_size=sp_size,
+ ep_size=ep_size,
+ zero_stage=stage,
+ enable_sequence_parallelism=sp_size > 1,
+ sequence_parallelism_mode="all_to_all" if sp_size > 1 else None,
+ overlap_communication=False,
+ initial_scale=1,
+ precision=precision,
+ find_unused_parameters=True,
+ )
+ dp_size = plugin.dp_size
+
+ booster = Booster(plugin=plugin)
+
+ assert pp_size <= NUM_LAYERS, "pp_size should be less than or equal to NUM_LAYERS"
+ config = MixtralConfig(
+ hidden_size=HIDDEN_SIZE_PER_HEAD * NUM_HEADS,
+ intermediate_size=HIDDEN_SIZE_PER_HEAD * NUM_HEADS * 2,
+ num_hidden_layers=NUM_LAYERS,
+ num_attention_heads=NUM_HEADS,
+ num_key_value_heads=NUM_HEADS,
+ num_local_experts=NUM_EXPERTS,
+ num_experts_per_tok=TOP_K,
+ attn_implementation="flash_attention_2",
+ )
+
+ # init model with the same seed
+ seed_all(10086)
+
+ torch_model = MixtralModel(config).to(dtype).cuda()
+ torch_optimizer = torch.optim.SGD(torch_model.parameters(), lr=1)
+
+ parallel_model = deepcopy(torch_model)
+ parallel_optimizer = torch.optim.SGD(parallel_model.parameters(), lr=1)
+ parallel_model, parallel_optimizer, _, _, _ = booster.boost(parallel_model, parallel_optimizer)
+
+ # create different input along dp axis
+ seed_all(1453 + rank)
+
+ torch_model.train()
+ parallel_model.train()
+ for _ in range(2):
+ # gen random input
+ input_embeddings = torch.rand(
+ NUM_BATCH, NUM_TOK_PER_BATCH, HIDDEN_SIZE_PER_HEAD * NUM_HEADS, requires_grad=True
+ ).cuda()
+ dist.all_reduce(
+ input_embeddings, group=plugin.pp_group
+ ) # pp inputs except the first stage doesn't matter, but need to be replicate for torch model check
+
+ dist.all_reduce(input_embeddings, group=plugin.tp_group) # tp group duplicate input
+ dist.all_reduce(input_embeddings, group=plugin.sp_group) # sp group duplicate input
+
+ # run the model with hybrid parallel
+ if booster.plugin.stage_manager is not None:
+ # for test with pp
+ data_iter = iter([{"inputs_embeds": input_embeddings}])
+ sharded_output = booster.execute_pipeline(
+ data_iter,
+ parallel_model,
+ lambda x, y: x.last_hidden_state.mean(),
+ parallel_optimizer,
+ return_loss=True,
+ return_outputs=True,
+ )
+ if booster.plugin.stage_manager.is_last_stage():
+ parallel_output = sharded_output["loss"]
+ else:
+ parallel_output = torch.tensor(12345.0, device="cuda")
+
+ # broadcast along pp axis
+ dist.broadcast(
+ parallel_output, src=dist.get_process_group_ranks(plugin.pp_group)[-1], group=plugin.pp_group
+ )
+ else:
+ # for test without pp
+ parallel_output = parallel_model(inputs_embeds=input_embeddings.to(dtype)).last_hidden_state.mean()
+ parallel_optimizer.backward(parallel_output)
+ parallel_optimizer.step()
+ parallel_optimizer.zero_grad()
+ dist.all_reduce(parallel_output, group=plugin.dp_group)
+
+ # ===================================================================================
+ # run normal model with all dp(different) inputs
+ all_inputs = [torch.empty_like(input_embeddings) for _ in range(dp_size)]
+ dist.all_gather(all_inputs, input_embeddings, group=plugin.dp_group)
+ torch_output_sum = 0
+ for input_data_ in all_inputs:
+ torch_output = torch_model(inputs_embeds=input_data_.to(dtype)).last_hidden_state.mean()
+ torch_output.backward()
+ torch_output_sum += torch_output.detach()
+ # avg dp grads follows zero optimizer
+ for p in torch_model.parameters():
+ if p.grad is not None:
+ p.grad /= dp_size
+ torch_optimizer.step()
+ torch_optimizer.zero_grad()
+
+ assert_loose_close(parallel_output, torch_output_sum, dtype=dtype)
+
+ # use checkpoint to load sharded zero model
+ model_dir = "./test_mixtral"
+ if rank == world_size - 1:
+ os.makedirs(model_dir, exist_ok=True)
+
+ dist.barrier()
+ booster.save_model(parallel_model, model_dir, shard=True)
+ dist.barrier()
+
+ saved_model = MixtralModel.from_pretrained(model_dir).cuda().to(dtype)
+ check_model_equal(torch_model, saved_model)
+ dist.barrier()
+
+ if rank == world_size - 1:
+ shutil.rmtree(model_dir)
+
+ print(f"rank {dist.get_rank()} test passed")
+
+
+def run_dist(rank, world_size, port):
+ colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
+ run_zero_with_original_model()
+
+
+@pytest.mark.dist
+@pytest.mark.parametrize("world_size", [4])
+@rerun_if_address_is_in_use()
+def test_mixtral(world_size):
+ spawn(run_dist, world_size)
+
+
+if __name__ == "__main__":
+ test_mixtral(world_size=4)
diff --git a/tests/test_shardformer/test_model/test_shard_qwen2.py b/tests/test_shardformer/test_model/test_shard_qwen2.py
index 166b31df9..c87415b75 100644
--- a/tests/test_shardformer/test_model/test_shard_qwen2.py
+++ b/tests/test_shardformer/test_model/test_shard_qwen2.py
@@ -135,6 +135,68 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn,
"precision": "fp16",
"initial_scale": 1,
},
+ { # Ulysess + Flash attention
+ "tp_size": 1,
+ "pp_size": 2,
+ "sp_size": 2,
+ "num_microbatches": 2,
+ "enable_sequence_parallelism": True,
+ "sequence_parallelism_mode": "all_to_all",
+ "enable_flash_attention": True,
+ "use_lazy_init": True,
+ "zero_stage": 1,
+ "precision": "fp16",
+ "initial_scale": 1,
+ },
+ {
+ "tp_size": 2,
+ "pp_size": 2,
+ "sp_size": 2,
+ "num_microbatches": 2,
+ "enable_sequence_parallelism": True,
+ "sequence_parallelism_mode": "split_gather",
+ "enable_flash_attention": True,
+ "use_lazy_init": True,
+ "zero_stage": 1,
+ "precision": "fp16",
+ "initial_scale": 1,
+ },
+ {
+ "tp_size": 2,
+ "pp_size": 2,
+ "sp_size": 2,
+ "num_microbatches": 2,
+ "enable_sequence_parallelism": True,
+ "sequence_parallelism_mode": "ring",
+ "enable_flash_attention": True,
+ "use_lazy_init": True,
+ "zero_stage": 1,
+ "precision": "fp16",
+ "initial_scale": 1,
+ },
+ {
+ "tp_size": 1,
+ "pp_size": 1,
+ "sp_size": 2,
+ "num_microbatches": 1,
+ "enable_sequence_parallelism": True,
+ "sequence_parallelism_mode": "all_to_all",
+ "use_lazy_init": True,
+ "zero_stage": 1,
+ "precision": "fp16",
+ "initial_scale": 1,
+ },
+ {
+ "tp_size": 4,
+ "pp_size": 1,
+ "num_microbatches": 1,
+ "enable_sequence_parallelism": True,
+ "sequence_parallelism_mode": "split_gather",
+ "enable_flash_attention": False,
+ "use_lazy_init": True,
+ "precision": "fp16",
+ "initial_scale": 1,
+ },
{
"tp_size": 1,
"pp_size": 2,
@@ -151,8 +213,11 @@ def run_qwen2_test(test_config):
sub_model_zoo = model_zoo.get_sub_registry("transformers_qwen2")
for name, (model_fn, data_gen_fn, output_transform_fn, loss_fn, _) in sub_model_zoo.items():
- check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, test_config)
-
+ try:
+ check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, test_config)
+ except Exception as e:
+ print(f"Failed config: {test_config}")
+ raise e
clear_layout_converter()
Randomizer.reset_index()
torch.cuda.empty_cache()
@@ -197,7 +262,11 @@ def run_qwen2_3d_test(test_config):
sub_model_zoo = model_zoo.get_sub_registry("transformers_qwen2")
for name, (model_fn, data_gen_fn, output_transform_fn, loss_fn, _) in sub_model_zoo.items():
- check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, test_config)
+ try:
+ check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, test_config)
+ except Exception as e:
+ print(f"Failed config: {test_config}")
+ raise e
clear_layout_converter()
Randomizer.reset_index()
diff --git a/tests/test_zero/test_low_level/test_grad_acc.py b/tests/test_zero/test_low_level/test_grad_acc.py
index ed12bb72d..94db70ca5 100644
--- a/tests/test_zero/test_low_level/test_grad_acc.py
+++ b/tests/test_zero/test_low_level/test_grad_acc.py
@@ -64,8 +64,12 @@ def exam_zero_1_2_grad_acc():
zero1_optimizer.step()
zero2_optimizer.step()
+ zero1_optimizer._force_wait_all_gather()
+ zero2_optimizer._force_wait_all_gather()
+
# check updated param
for z1p, z2p in zip(zero1_model.parameters(), zero2_model.parameters()):
+ assert not hasattr(z1p, "_all_gather_handle")
assert torch.equal(z1p.data, z2p.data)
diff --git a/tests/test_zero/test_low_level/test_zero1_2.py b/tests/test_zero/test_low_level/test_zero1_2.py
index cda51c4ef..368c782fe 100644
--- a/tests/test_zero/test_low_level/test_zero1_2.py
+++ b/tests/test_zero/test_low_level/test_zero1_2.py
@@ -190,6 +190,8 @@ def exam_zero_1_torch_ddp(world_size, dtype: torch.dtype, master_weights: bool):
# torch ddp step
torch_optimizer.step()
+ zero_optimizer._force_wait_all_gather()
+
# check updated param
for (n, p), z1p in zip(torch_model.named_parameters(), zero_model.parameters()):
loose_close(p, z1p, dtype=dtype)
diff --git a/version.txt b/version.txt
index 1d0ba9ea1..2b7c5ae01 100644
--- a/version.txt
+++ b/version.txt
@@ -1 +1 @@
-0.4.0
+0.4.2