tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

wip for tp

Browse Source
This commit is contained in:
Wing Lian
2024-08-23 10:57:57 -04:00
parent 198f7cd893
commit 756a34f0fe
1 changed files with 16 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
src/axolotl/core/trainer_builder.py
View File

@@ -1252,8 +1252,9 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
"cuda", (dp_size, tp_size), mesh_dim_names=("dp", "tp") "cuda", (dp_size, tp_size), mesh_dim_names=("dp", "tp")
) )
dp_mesh = device_mesh["dp"] dp_mesh = device_mesh["dp"]
tp_mesh = device_mesh["tp"]
training_arguments_kwargs["fsdp_config"]["device_mesh"] = dp_mesh training_arguments_kwargs["fsdp_config"]["device_mesh"] = dp_mesh
self.parallelize_model(device_mesh) self.parallelize_model(tp_mesh)
if self.cfg.adapter == "qlora": if self.cfg.adapter == "qlora":
training_arguments_kwargs["qlora"] = True training_arguments_kwargs["qlora"] = True
@@ -1626,7 +1627,7 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
return trainer return trainer
def parallelize_model(self, device_mesh, loss_parallel=True): def parallelize_model(self, device_mesh, loss_parallel=False):
# FIXME hardcoded for llama # FIXME hardcoded for llama
tp_mesh = device_mesh["tp"] tp_mesh = device_mesh["tp"]
@@ -1641,7 +1642,6 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
), ),
}, },
) )
parallelize_module( parallelize_module(
self.model.model, self.model.model,
tp_mesh, tp_mesh,
@@ -1654,12 +1654,12 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
}, },
) )
for _, transformer_block in self.model.model.layers.items(): for _, transformer_block in enumerate(self.model.model.layers):
layer_plan = { layer_plan = {
"input_layernorm": SequenceParallel(), "input_layernorm": SequenceParallel(),
"self_attn": PrepareModuleInput( "self_attn": PrepareModuleInput(
input_layouts=(Shard(1), None), input_layouts=(Shard(1),),
desired_input_layouts=(Replicate(), None), desired_input_layouts=(Replicate()),
), ),
"self_attn.q_proj": ColwiseParallel(), "self_attn.q_proj": ColwiseParallel(),
"self_attn.k_proj": ColwiseParallel(), "self_attn.k_proj": ColwiseParallel(),
@@ -1671,9 +1671,17 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
desired_input_layouts=(Replicate(),), desired_input_layouts=(Replicate(),),
), ),
"mlp.gate_proj": ColwiseParallel(), "mlp.gate_proj": ColwiseParallel(),
"mlp.up_proj": RowwiseParallel(output_layouts=Shard(1)), "mlp.up_proj": ColwiseParallel(),
"mlp.down_proj": ColwiseParallel(), "mlp.down_proj": RowwiseParallel(output_layouts=Shard(1)),
} }
self_attn = transformer_block.self_attn
self_attn.num_heads = self_attn.num_heads // tp_mesh.size()
self_attn.num_key_value_heads = (
self_attn.num_key_value_heads // tp_mesh.size()
)
# TODO need to fix self_attn.rotary_emb
parallelize_module( parallelize_module(
transformer_block, transformer_block,
tp_mesh, tp_mesh,