Skip to content

Commit 88b4f94

Browse files
alay2shahalay2shah
authored
MoE Support (#7)
* transformers bump * config updates * config readme * disable use cache for act checkpointing * Moe DPO + DS configs * pid added to triton cache * MoE SFT + DS * fsdp dpo * fsdp sft * add grad accum steps to user config * dtype change * util for moe detection * ray tmp * fsdp moe config * adding moe to sft run * adding moe to dpo run * moe configs * rm docstring
1 parent 1af6122 commit 88b4f94

15 files changed

Lines changed: 1327 additions & 1028 deletions

‎config.py‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -42,11 +42,15 @@
4242
Default SFT: TrainingConfig.DEFAULT_SFT
4343
Default DPO: TrainingConfig.DEFAULT_DPO
4444
Default VLM SFT: TrainingConfig.DEFAULT_VLM_SFT
45+
MoE SFT: TrainingConfig.MOE_SFT
46+
MoE DPO: TrainingConfig.MOE_DPO
4547
4648
No LORA (full finetuning): PeftConfig.NO_LORA
4749
Default LORA: PeftConfig.DEFAULT_LORA
4850
High R LoRA: PeftConfig.HIGH_R_LORA
4951
VLM LoRA: PeftConfig.DEFAULT_VLM_LORA
52+
MoE LoRA: PeftConfig.MOE_LORA
53+
MoE High R LoRA: PeftConfig.MOE_LORA_HIGH_R
5054
5155
Args:
5256
output_dir: Output directory for training artifacts
@@ -59,10 +63,9 @@
5963
"output_dir": None,
6064
"num_train_epochs": None,
6165
"per_device_train_batch_size": None,
66+
"gradient_accumulation_steps": None,
6267
"learning_rate": None,
6368
}
64-
training_config = TrainingConfig.DEFAULT_SFT.override(**user_config)
65-
peft_config = PeftConfig.DEFAULT_LORA
6669

6770

6871
#################################
@@ -74,18 +77,18 @@
7477
7578
Args:
7679
job_name: Unique identifier for the training job (alphanumeric, hyphens, underscores only)
77-
model_name: Model to fine-tune - default: "LFM2-1.2B". Try also "LFM2-700M" or "LFM2-350M"
78-
training_type: "sft" or "dpo"
80+
model_name: Model to fine-tune - default: "LFM2-1.2B". Try also "LFM2-700M", "LFM2-350M", or "LFM2-8B-A1B" (MoE)
81+
training_type: "sft", "dpo", or "vlm_sft"
7982
dataset: Dataset to use for training, defined in step 1
8083
training_config: Training configuration that matches training_type, defined in step 2
8184
peft_config: PEFT configuration for parameter-efficient fine-tuning, defined in step 2
8285
"""
8386

8487
JOB_CONFIG = JobConfig(
8588
job_name="my_job_name",
86-
model_name="LFM2-1.2B",
89+
model_name="LFM2-8B-A1B",
8790
training_type="sft",
8891
dataset=example_sft_dataset,
89-
training_config=training_config,
90-
peft_config=peft_config,
92+
training_config=TrainingConfig.MOE_SFT.override(**user_config),
93+
peft_config=PeftConfig.NO_LORA,
9194
)

‎pyproject.toml‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ dependencies = [
1313
"deepspeed>=0.17.1",
1414
"ray[train]>=2.47.1",
1515
"torch>=2.7.1",
16-
"transformers>=4.55.0",
16+
"transformers @ git+https://github.com/huggingface/transformers.git@0c9a72e4576fe4c84077f066e585129c97bfd4e6",
1717
"trl>=0.18.2",
1818
"rich>=14.1.0",
1919
"pillow>=11.3.0",
@@ -28,6 +28,9 @@ leap-finetune = "leap_finetune:main"
2828
requires = ["hatchling"]
2929
build-backend = "hatchling.build"
3030

31+
[tool.hatch.metadata]
32+
allow-direct-references = true
33+
3134
[dependency-groups]
3235
dev = [
3336
"pre-commit>=4.2.0",

‎src/leap_finetune/configs/__init__.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
from enum import Enum
22
from .dpo_configs import (
33
DEFAULT_DPO_CONFIG,
4+
MOE_DPO_CONFIG,
45
)
56
from .sft_configs import (
67
DEFAULT_SFT_CONFIG,
8+
MOE_SFT_CONFIG,
79
)
810
from .vlm_sft_config import (
911
DEFAULT_VLM_SFT_CONFIG,
@@ -12,13 +14,17 @@
1214
LFM2_LORA_DEFAULT_CONFIG,
1315
LFM2_LORA_HIGH_R_CONFIG,
1416
LFM2_VLM_LORA_CONFIG,
17+
LFM2_MOE_LORA_CONFIG,
18+
LFM2_MOE_LORA_HIGH_R_CONFIG,
1519
)
1620

1721

1822
class TrainingConfig(Enum):
1923
DEFAULT_SFT = DEFAULT_SFT_CONFIG
2024
DEFAULT_DPO = DEFAULT_DPO_CONFIG
2125
DEFAULT_VLM_SFT = DEFAULT_VLM_SFT_CONFIG
26+
MOE_SFT = MOE_SFT_CONFIG
27+
MOE_DPO = MOE_DPO_CONFIG
2228

2329
def override(self, **overrides):
2430
"""Create a custom TrainingConfig with overrides"""
@@ -38,3 +44,5 @@ class PeftConfig(Enum):
3844
DEFAULT_LORA = LFM2_LORA_DEFAULT_CONFIG
3945
HIGH_R_LORA = LFM2_LORA_HIGH_R_CONFIG
4046
DEFAULT_VLM_LORA = LFM2_VLM_LORA_CONFIG
47+
MOE_LORA = LFM2_MOE_LORA_CONFIG
48+
MOE_LORA_HIGH_R = LFM2_MOE_LORA_HIGH_R_CONFIG

‎src/leap_finetune/configs/distributed_configs.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
"""Shared distributed training configurations."""
2+
3+
########################
4+
# FSDP CONFIGS #
5+
########################
6+
7+
MOE_FSDP_CONFIG = {
8+
"fsdp": ["shard_grad_op", "auto_wrap"],
9+
"fsdp_config": {
10+
"transformer_layer_cls_to_wrap": "Lfm2MoeDecoderLayer",
11+
"backward_prefetch": "backward_pre",
12+
"sync_module_states": True,
13+
"use_orig_params": False,
14+
},
15+
}

‎src/leap_finetune/configs/dpo_configs.py‎

Lines changed: 54 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,11 @@
22

33

44
########################
5-
# DEEPSEED CONFIGS #
5+
# DEEPSPEED CONFIGS #
66
########################
77

88

9-
DEEPSEED_CONFIG = {
9+
DEEPSPEED_CONFIG = {
1010
"zero_optimization": {
1111
"stage": 2,
1212
"overlap_comm": True,
@@ -27,6 +27,27 @@
2727
}
2828

2929

30+
MOE_DEEPSPEED_CONFIG = {
31+
"zero_optimization": {
32+
"stage": 0,
33+
"overlap_comm": True,
34+
},
35+
"train_batch_size": "auto",
36+
"train_micro_batch_size_per_gpu": "auto",
37+
"gradient_clipping": "auto",
38+
"gradient_accumulation_steps": "auto",
39+
"bf16": {"enabled": "auto"},
40+
"activation_checkpointing": {
41+
"partition_activations": False,
42+
"cpu_checkpointing": False,
43+
"contiguous_memory_optimization": False,
44+
"number_checkpoints": None,
45+
"synchronize_checkpoint_boundary": False,
46+
"profile": False,
47+
},
48+
}
49+
50+
3051
########################
3152
# DPO CONFIGS #
3253
########################
@@ -42,9 +63,39 @@
4263
"beta": 0.1,
4364
"loss_type": "sigmoid",
4465
"logging_steps": 10,
66+
"logging_first_step": True,
4567
"save_strategy": "epoch",
4668
"eval_strategy": "epoch",
4769
"load_best_model_at_end": True,
4870
"ddp_find_unused_parameters": False,
49-
"deepspeed": DEEPSEED_CONFIG,
71+
"deepspeed": DEEPSPEED_CONFIG,
72+
}
73+
74+
75+
########################
76+
# MOE DPO CONFIGS #
77+
########################
78+
79+
# Base MoE DPO config - distributed strategy is applied automatically in runner
80+
# based on PEFT presence: DeepSpeed for LoRA, FSDP for full fine-tuning
81+
MOE_DPO_CONFIG = {
82+
"training_type": "dpo",
83+
"output_dir": DPO_OUTPUT_PATH,
84+
"num_train_epochs": 2, # MoE models typically need fewer epochs
85+
"per_device_train_batch_size": 2, # MoE models are larger, use smaller batch size
86+
"learning_rate": 1e-6,
87+
"lr_scheduler_type": "linear",
88+
"beta": 0.1,
89+
"loss_type": "sigmoid",
90+
"logging_steps": 10,
91+
"logging_first_step": True,
92+
"save_strategy": "epoch",
93+
"eval_strategy": "epoch",
94+
"load_best_model_at_end": True,
95+
"max_grad_norm": 1.0,
96+
"bf16": True,
97+
# Distributed strategy will be set automatically:
98+
# - With PEFT: uses MOE_DEEPSPEED_CONFIG
99+
# - Without PEFT: uses FSDP_CONFIG
100+
"deepspeed": MOE_DEEPSPEED_CONFIG,
50101
}

‎src/leap_finetune/configs/peft_configs.py‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,3 +37,28 @@
3737
bias="none",
3838
target_modules=LFM_MODULES + VISION_TOWER_MODULES + MULTI_MODAL_PROJECTOR_MODULES,
3939
)
40+
41+
42+
########################
43+
# MOE CONFIGS #
44+
########################
45+
46+
LFM2_MOE_LORA_CONFIG = LoraConfig(
47+
task_type=TaskType.CAUSAL_LM,
48+
inference_mode=False,
49+
r=8,
50+
lora_alpha=16,
51+
lora_dropout=0.1,
52+
bias="none",
53+
target_modules="all-linear", # Target all linear layers in MoE architecture
54+
)
55+
56+
LFM2_MOE_LORA_HIGH_R_CONFIG = LoraConfig(
57+
task_type=TaskType.CAUSAL_LM,
58+
inference_mode=False,
59+
r=16,
60+
lora_alpha=32,
61+
lora_dropout=0.1,
62+
bias="none",
63+
target_modules="all-linear", # Target all linear layers in MoE architecture
64+
)

‎src/leap_finetune/configs/sft_configs.py‎

Lines changed: 62 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,11 @@
22

33

44
########################
5-
# DEEPSEED CONFIGS #
5+
# DEEPSPEED CONFIGS #
66
########################
77

88

9-
DEEPSEED_CONFIG = {
9+
DEEPSPEED_CONFIG = {
1010
"zero_optimization": {
1111
"stage": 2,
1212
"overlap_comm": True,
@@ -36,6 +36,36 @@
3636
}
3737

3838

39+
MOE_DEEPSPEED_CONFIG = {
40+
"zero_optimization": {
41+
"stage": 0,
42+
"overlap_comm": True,
43+
},
44+
"train_batch_size": "auto",
45+
"train_micro_batch_size_per_gpu": "auto",
46+
"gradient_clipping": "auto",
47+
"gradient_accumulation_steps": "auto",
48+
"optimizer": {
49+
"type": "AdamW",
50+
"params": {
51+
"lr": "auto", # Uses learning_rate from training config
52+
"betas": "auto", # DEFAULT: (0.9, 0.999)
53+
"eps": "auto", # DEFAULT: 1e-8
54+
"weight_decay": "auto", # DEFAULT: 0.01
55+
},
56+
},
57+
"fp16": {"enabled": "auto"},
58+
"activation_checkpointing": {
59+
"partition_activations": False,
60+
"cpu_checkpointing": False,
61+
"contiguous_memory_optimization": False,
62+
"number_checkpoints": None,
63+
"synchronize_checkpoint_boundary": False,
64+
"profile": False,
65+
},
66+
}
67+
68+
3969
########################
4070
# SFT CONFIGS #
4171
########################
@@ -55,5 +85,34 @@
5585
"eval_strategy": "epoch",
5686
"load_best_model_at_end": True,
5787
"ddp_find_unused_parameters": False,
58-
"deepspeed": DEEPSEED_CONFIG,
88+
"deepspeed": DEEPSPEED_CONFIG,
89+
}
90+
91+
92+
########################
93+
# MOE SFT CONFIGS #
94+
########################
95+
96+
# Base MoE SFT config - distributed strategy is applied automatically in runner
97+
# based on PEFT presence: DeepSpeed for LoRA, FSDP for full fine-tuning
98+
MOE_SFT_CONFIG = {
99+
"training_type": "sft",
100+
"output_dir": SFT_OUTPUT_PATH,
101+
"num_train_epochs": 2, # MoE models typically need fewer epochs
102+
"per_device_train_batch_size": 2, # Reduced to save memory
103+
"gradient_accumulation_steps": 1, # Set to 1 to match Accelerate config for testing
104+
"learning_rate": 5e-5,
105+
"lr_scheduler_type": "linear",
106+
"warmup_steps": 100,
107+
"warmup_ratio": 0.2,
108+
"logging_steps": 10,
109+
"save_strategy": "epoch",
110+
"eval_strategy": "epoch",
111+
"load_best_model_at_end": True,
112+
"max_grad_norm": 1.0,
113+
"bf16": True,
114+
# Distributed strategy will be set automatically:
115+
# - With PEFT: uses MOE_DEEPSPEED_CONFIG
116+
# - Without PEFT: uses FSDP_CONFIG
117+
"deepspeed": MOE_DEEPSPEED_CONFIG,
59118
}

‎src/leap_finetune/configs/vlm_sft_config.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
11
from leap_finetune.utils.constants import SFT_OUTPUT_PATH
22

33
########################
4-
# DEEPSEED CONFIGS #
4+
# DEEPSPEED CONFIGS #
55
########################
66

77

8-
DEEPSEED_CONFIG = {
8+
DEEPSPEED_CONFIG = {
99
"zero_optimization": {
1010
"stage": 2,
1111
"overlap_comm": True,
@@ -55,5 +55,5 @@
5555
"load_best_model_at_end": True,
5656
"gradient_checkpointing": True,
5757
"dataset_kwargs": {"skip_prepare_dataset": True},
58-
"deepspeed": DEEPSEED_CONFIG,
58+
"deepspeed": DEEPSPEED_CONFIG,
5959
}

‎src/leap_finetune/trainer.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,9 +34,10 @@ def ray_trainer(job_config: dict) -> None:
3434
if not cuda.is_available():
3535
raise ValueError("No GPU available for training")
3636

37+
ray_temp_dir = os.path.expanduser("~/ray_temp")
38+
os.makedirs(ray_temp_dir, exist_ok=True)
39+
3740
if not ray.is_initialized():
38-
ray_temp_dir = os.path.expanduser("~/ray_temp")
39-
os.makedirs(ray_temp_dir, exist_ok=True)
4041
ray.init(
4142
runtime_env=RuntimeEnv(
4243
working_dir=str(RUNTIME_DIR),
@@ -45,6 +46,7 @@ def ray_trainer(job_config: dict) -> None:
4546
"TEMP": ray_temp_dir,
4647
"TMP": ray_temp_dir,
4748
"NCCL_IB_DISABLE": "1",
49+
"NCCL_P2P_DISABLE": "1", # Added to force CPU level communication during synchronization
4850
"TORCH_NCCL_ASYNC_ERROR_HANDLING": "1",
4951
"NCCL_SOCKET_IFNAME": "lo",
5052
"TORCH_NCCL_BLOCKING_WAIT": "1",

0 commit comments

Comments
 (0)