axolotl : resume_from_checkpoint 시, save checkpoint 및 resume train error

axolotl로 qdora 학습을 진행 하는 중 마주친 에러이다.

Error invalid argument at line 396 in file /src/csrc/pythonInterface.cpp

해당 에러는 axolotl config 중  optimizer: paged_adamw_32bit로 설정하여 발생하는 에러인데, 이 인자를 사용하면, 실제로 gpu 메모리가 10GB 정도 적게 사용되며, 학습을 할 수 있었다.

 

하지만, saves_per_epoch 인자와 호환이 안되는 문제가 있어서, save_steps로 steps마다 저장을 해주었지만, 마지막 iter의 checkpoint에 저장하는데에 오류가 발생했다. 

 

문제 해결은 다음과 같이, 해당 checkpoint의 optimizer.pt를 만져준 뒤, optimizer: torch_adamw로 변경하여 resume_from_checkpoint를 활성화 하면, 정상적으로 checkpoint로 부터 이어서 학습이 가능했다.

import os
import torch

if __name__ == "__main__":
    os.rename("optimizer.pt", "optimizer.paged.pt")
    state_dict = torch.load("optimizer.paged.pt")
    for i in state_dict["state"].keys():
        exp_avg = state_dict["state"][i].pop("state1")
        exp_avg_sq = state_dict["state"][i].pop("state2")
        state_dict["state"][i]["exp_avg"] = exp_avg
        state_dict["state"][i]["exp_avg_sq"] = exp_avg_sq

    for i in range(len(state_dict["param_groups"])):
        state_dict["param_groups"][i]["amsgrad"] = False
        state_dict["param_groups"][i]["foreach"] = None
        state_dict["param_groups"][i]["maximize"] = False
        state_dict["param_groups"][i]["capturable"] = False
        state_dict["param_groups"][i]["differentiable"] = False
        state_dict["param_groups"][i]["fused"] = None

    torch.save(state_dict, "optimizer.pt")

 

saves_per_epoch: 1 설정 시, checkpoint 저장시 NCCL 에러가 발생하는데, 이 것도 optimizer 문제인지는 실험을 해봐야 알 것 같다.