Axotl & deepspeed zero3: 모델 레이어 불러올때 torch.Size([0]) 해결법

진짜 그냥 이걸 해결했긴 했는데, 진짜 머리가 찡하다.

 

그냥 화가 무척 난다.

 

왜 화가나느냐?

 

마이크로 소프트 이놈들은 deepspeed를 윈도우처럼 만들어놨다.

 

업데이트도 느리고 github 레포에서 세달전에 올라온 버그 리포트도 아직까지 반영이 안되어있다.

 

자 그럼 이번에 어떤주제냐면, LLM full-finetuning할때, 내가 몇개의 레이어는 얼리고 몇개는 풀고싶고 이렇게 튜닝하고 싶은 경우가 정말 많다. 

 

그럴때 자주 쓰이는 것은 아래와 같다.

 

param.requires_grad = False 혹은 param.requires_grad = True

 

 

위의 두개가 무엇이냐 하면은, 그 가중치에 gradient update를 할지 안할지 결정해주는 것이다.

 

그런데 이제 Axotl라이브러리 활용하셔 LLM을 full-fine tuning을 multi-GPU로 하고자 하면서 레이어를 얼려주고 풀어주고자 하면 다음과 같은 문제점이 발생한다. 

 

model.embed_tokens.weight
torch.Size([32000, 4096])


model.layers.8.self_attn.q_proj.weight
torch.Size([0])

 

나는 embed layer와 8번 레이어의 가중치를 업데이트 하고 나머지는 묶어주고자 하였는데 lm_head나 embed_tokens 텐서는 잘 불러와 지는데 나머지 레이어의 텐서는 인스턴스화 되지 않고 그냥 불러와 지지 않는 문제가 발생한다.

 

자 여기서 나는 수많은 깃허브 이슈들을 읽으면서 내린 결론은 다음과 같다.

 

"Deepseed의 zero3 문제"

Deepseed는 분산학습에서 매우 메모리를 효율적으로 설계하여 훈련시킬 수 있는 툴이다. 다만, 몇가지 그리고 많은 문제점들이 있지만...

그래도 해결되면 매우 간편하고 효율적이다.

 

그런데, 여기서 zero initialize를 하면서 명시적으로 인스턴스화 해주지 않고 불러오기 때문에, 모델 내부 레이어는 명시적으로 인스턴스화 되지 않는다.

 

그렇기에 내가 인스턴스화 되지 않는 레이어를 풀어주고자 해도 절대로 안풀어지는 법...

 

그럼 어떻게 해결해야 하는가?

 

Axotl의 freeze.py를 보면 다음과 같이 코드가 작성되어 있다.

 

for param in model.parameters():

    param.requires_grad = False



for name, param in model.named_parameters():

    if any(pattern.match(name) for pattern in compiled_patterns):

        if is_main_process():

               LOG.debug(f"unfreezing {name}")

        param.requires_grad = True

 

 

즉 모든 레이어를 얼리고, 내가 풀고자 하는 레이어만 쓱 푸는것이다.

 

근데, 애초에 모델을 불러오는데에 있어서 모든 레이어는 required_grad가 True로 되어 있다.

 

그러면, 굳이 얼리고 풀면서 생기는 저런 이슈를 마주쳐야 하는가?

전혀 아니다. 

 

인스턴스화 되지 않은것에 대해서 required_grad=True가 문제가 되지 사실 required_grad=False는 문제가 되지 않는다.

 

그러면 어떻게 해결해야 하는가?

 

# for param in model.parameters():
  # param.requires_grad = False

for name, param in model.named_parameters():
  if any(pattern.match(name) for pattern in compiled_patterns):
    if is_main_process():
        LOG.debug(f"unfreezing {name}")
   # param.requires_grad = True
    
    pass
    
  else:
    param.requires_grad = False

 

방법은 발상의 거꾸로 전환이다.

그냥 처음부터 얼리지말고, 내가 선택하고 싶은건 pass하고 내가 선택하지 않은 레이어만 얼려주는 것이다.

 

하... Axotl도 답답하다..... 하... 내 13시간.....