github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-08 20:40:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

[exampe] update llama example (#5626)

Browse Source 
* [plugin] support dp inside for hybriad parallel

* [example] update llama benchmark

* [example] update llama benchmark

* [example] update llama readme

* [example] update llama readme
This commit is contained in:
Hongxin Liu
2024-04-23 14:12:20 +08:00
committed by GitHub
parent 862fbaaa62
commit 4de4e31818
8 changed files with 72 additions and 783 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
colossalai/booster/plugin/gemini_plugin.py
View File

@@ -424,6 +424,7 @@ class GeminiPlugin(DPPluginBase):
)
self.extra_dp_group = self.pg_mesh.get_group_along_axis(DP_AXIS) if self.extra_dp_size > 1 else None
self.tp_group = self.pg_mesh.get_group_along_axis(TP_AXIS) if self.tp_size > 1 else None
self.dp_size = self.zero_size * self.extra_dp_size
self.shard_config = ShardConfig(
tensor_parallel_process_group=self.tp_group,