github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-10-10 08:44:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

[builder] raise Error when CUDA_HOME is not set (#2213)

Browse Source
This commit is contained in:
Jiarui Fang
2022-12-28 16:07:08 +08:00
committed by GitHub
parent 78a89d9b41
commit 7675792100
7 changed files with 19 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
examples/language/gpt/train_gpt_demo.py
View File

@@ -218,7 +218,7 @@ def main():
model = gemini_zero_dpp(model, pg, args.placement)
# build highly optimized cpu optimizer
optimizer = GeminiAdamOptimizer(model, lr=1e-3, initial_scale=2**5)
optimizer = GeminiAdamOptimizer(model, lr=1e-3, initial_scale=2**5, gpu_margin_mem_ratio=0.6)
logger.info(get_mem_info(prefix='After init optim, '), ranks=[0])
else:
model = model_builder(args.model_type)(checkpoint=True).cuda()