github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[pre-commit.ci] auto fixes from pre-commit.com hooks 

for more information, see https://pre-commit.ci
pull/5722/head
pre-commit-ci[bot] 2024-05-16 05:20:00 +00:00
parent 4148ceed9f
commit 5bedea6e10
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -329,7 +329,7 @@ class GeminiPlugin(DPPluginBase):
chunk_init_device: Optional[torch.device] = None, chunk_init_device: Optional[torch.device] = None,
placement_policy: str = "static", placement_policy: str = "static",
enable_gradient_accumulation: bool = False, enable_gradient_accumulation: bool = False,
max_prefetch:int = 0, max_prefetch: int = 0,
shard_param_frac: float = 1.0, # only for static placement shard_param_frac: float = 1.0, # only for static placement
offload_optim_frac: float = 0.0, # only for static placement offload_optim_frac: float = 0.0, # only for static placement
offload_param_frac: float = 0.0, # only for static placement offload_param_frac: float = 0.0, # only for static placement