search

OpenLMLab / LOMO

LOMO: LOw-Memory Optimization
MIT License
978 stars 68 forks source link

LOMO是否支持bfloat16模型的训练？ #46

Closed Wangyupei closed 1 year ago

Wangyupei commented 1 year ago

非常感谢您的代码。对于float16类型的Llama-7B模型，代码运行正常。注意到代码支持bfloat16类型，我们把Llama模型转换为bfloat16类型，目前代码可以正常加载bfloat模型，也可以正常前向传播并计算loss值，但是反向传播时却报错： 代码运行至：self.optimizer.grad_norm(loss)，报错：AttributeError: 'NoneType' object has no attribute '_has_inf_or_nan'

lomo_bf_1

请问是什么原因？

KaiLv69 commented 1 year ago

您好，需要把ds_config中的fp16置为False，bf16置为True。还可以试试在代码最开始加上torch.set_default_dtype(torch.bfloat16)。

Wangyupei commented 1 year ago

感谢您的回复。 您提到的两处参数设置，之前我们在代码里已经修改了，但是还是出现上面的问题：训练时，bf模型能加载和前向传播，但反向传播就报错。 您之前测试bf16模型，只需要修改下这两处，就可以正常运行吗？

KaiLv69 commented 1 year ago

您好。 在代码最开始加上torch.set_default_dtype(torch.bfloat16)能解决您的问题吗？train_lomo_lora.py的代码中就使用了bf16.

Wangyupei commented 1 year ago

谢谢。 train_lomo_lora.py是可以正常训练bfloat16，但是train_lomo.py训练blfoat16会出错。 两处都已经修改了：ds_config中的fp16置为False，bf16置为True，也在train_lomo.py开始训练处添加torch.set_default_dtype(torch.bfloat16)。 但是，还是上面的报错：AttributeError: 'NoneType' object has no attribute '_has_inf_or_nan'，报错的代码是：> /output/LOMO/src/lomo.py(187)grad_norm() -> loss.backward(retain_graph=True) 请问是train_lomo.py不支持bfloat16训练吗？还是代码里需要修改？

KaiLv69 commented 1 year ago

可以尝试把 https://github.com/OpenLMLab/LOMO/blob/24cde8e91feac437809bf7790f4727623dce6a76/src/lomo_trainer.py#L88 里的model改为self.model。

Wangyupei commented 1 year ago

谢谢