Configuration for Megatron parallelism. The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config. Args: param_offload (bool): Whether to offload parameters to CPU. grad_offload (bool): Whether to offload gradients to CPU. o
| 78 | |
| 79 | @dataclass |
| 80 | class McoreEngineConfig(EngineConfig): |
| 81 | """Configuration for Megatron parallelism. |
| 82 | |
| 83 | The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config. |
| 84 | |
| 85 | Args: |
| 86 | param_offload (bool): Whether to offload parameters to CPU. |
| 87 | grad_offload (bool): Whether to offload gradients to CPU. |
| 88 | optimizer_offload (bool): Whether to offload optimizer states to CPU. |
| 89 | tensor_model_parallel_size (int): Tensor model parallel size. |
| 90 | expert_model_parallel_size (int): Expert model parallel size for MoE models. |
| 91 | expert_tensor_parallel_size (Optional[int]): Expert tensor parallel size for MoE models. |
| 92 | pipeline_model_parallel_size (int): Pipeline model parallel size. |
| 93 | virtual_pipeline_model_parallel_size (Optional[int]): Virtual pipeline model parallel size |
| 94 | for interleaved scheduling. |
| 95 | context_parallel_size (int): Context parallel size for long sequences. |
| 96 | sequence_parallel (bool): Whether to enable sequence parallelism. |
| 97 | use_distributed_optimizer (bool): Whether to use distributed optimizer. |
| 98 | use_dist_checkpointing (bool): Whether to use distributed checkpointing. |
| 99 | dist_checkpointing_path (Optional[str]): Path for distributed checkpointing. |
| 100 | seed (int): Random seed for reproducibility. |
| 101 | override_ddp_config (dict[str, Any]): Override configuration for DDP. |
| 102 | override_transformer_config (dict[str, Any]): Override configuration for transformer. |
| 103 | use_mbridge (bool): Whether to use MBridge for communication. |
| 104 | dtype (str): Mixed precision training param dtype, default "bfloat16" |
| 105 | """ |
| 106 | |
| 107 | # sequence_parallel is not listed as a frozen field for auto-correction purpose |
| 108 | _mutable_fields = EngineConfig._mutable_fields | {"sequence_parallel"} |
| 109 | # mcore parallelism |
| 110 | tensor_model_parallel_size: int = 1 |
| 111 | expert_model_parallel_size: int = 1 |
| 112 | expert_tensor_parallel_size: Optional[int] = None |
| 113 | pipeline_model_parallel_size: int = 1 |
| 114 | virtual_pipeline_model_parallel_size: Optional[int] = None |
| 115 | context_parallel_size: int = 1 |
| 116 | sequence_parallel: bool = True |
| 117 | use_distributed_optimizer: bool = True |
| 118 | use_dist_checkpointing: bool = False |
| 119 | dist_checkpointing_path: Optional[str] = None |
| 120 | dist_checkpointing_prefix: str = "" |
| 121 | override_ddp_config: dict[str, Any] = field(default_factory=dict) |
| 122 | override_transformer_config: dict[str, Any] = field(default_factory=dict) |
| 123 | override_mcore_model_config: dict[str, Any] = field(default_factory=dict) |
| 124 | use_mbridge: bool = True |
| 125 | vanilla_mbridge: bool = True |
| 126 | strategy: str = "megatron" |
| 127 | |
| 128 | def __post_init__(self) -> None: |
| 129 | super().__post_init__() |
| 130 | """config validation logics go here""" |
| 131 | assert self.strategy == "megatron" |
| 132 | assert self.dtype in ["bfloat16", "float16"], f"dtype {self.dtype} not supported" |
| 133 | if self.tensor_model_parallel_size == 1: |
| 134 | warnings.warn("set sequence parallel to false as TP size is 1", stacklevel=2) |
| 135 | self.sequence_parallel = False |
| 136 | |
| 137 |
no outgoing calls