Writer: Joonwon Jang

Contact: [email protected]

**** LLM Fine-Tuning시 Optimizer가 차지하는 VRAM의 비율이 높다고 하고 특히나 Learning Rate의 설정이 중요한 것을 깨닫는 요즘입니다🤣, 오늘 그 2가지가 왜 LLM 학습에 중요한지 제가 이해한 바를 적어보겠습니다! ****

⇒ 모델의 각 Parameter와 동일한 크기를 가지는 First Order Momentum과 Second Order Momentum을 사용하는 순간 모델 전체 Parameter 크기의 2배에 해당하는 추가 메모리를 VRAM에 로드해야 됩니다.

⇒ Gradient clipping을 하지 않았다고 가정했을 시 (물론 설정해도 마찬가지 입니다.), Gradient가 너무 커지면 학습률 Learning Rate가 적정하다고 판단된 값이라도 모델 Parameter나 Activation 값이 NaN(Not a Number)으로 수렴할 수 있습니다.

Optimizer’s Configuration

⇒ Mistral 7B를 AdamW를 활용해 1 epoch 학습하고 나면 아래와 같은 optimizer state가 출력됩니다.

Optimizer state at beginning of epoch 1:
{'state': {0: {'step': tensor([584.], device='cuda:2'), 'exp_avg': tensor([ 2.2361e-07, -1.4407e-07,  1.0716e-06,  ..., -3.3894e-07,
         0.0000e+00,  0.0000e+00], device='cuda:2'), 'exp_avg_sq': tensor([2.2285e-12, 3.9868e-12, 5.2110e-12,  ..., 5.1709e-12, 0.0000e+00,
         
....
0.0000e+00], device='cuda:2')}}, 'param_groups': [{'lr': 2e-06, 'betas': (0.9, 0.999), 'eps': 1e-08, 'weight_decay': 0.01, 'amsgrad': False, 'foreach': None, 'maximize': False, 'capturable': False, 'differentiable': False, 'fused': True, 'params': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34]}]}

• Optimizer의 전체 구조는 아래와 같은데요.

{
  'state': { ... },
  'param_groups': [ ... ]
}

• state: 모델의 각 parameter에 대한 optimizer의 상태를 저장하는 딕셔너리입니다. 키는 parameter의 index이고, value은 해당 parameter의 state 딕셔너리입니다.
• param_groups: optimizer의 하이퍼파라미터와 parameter 그룹에 대한 정보를 담고 있는 리스트입니다. 일반적으로 하나의 딕셔너리를 포함하며, 그 안에 lr, momentum term 등과 함께 파라미터의 인덱스 리스트(params)가 포함됩니다.

State’s Configuration

0: {
  'step': tensor([584.], device='cuda:2'),
  'exp_avg': tensor([...], device='cuda:2'),
  'exp_avg_sq': tensor([...], device='cuda:2')
},
1: { ... },

state 는 아래와 같은 필드로 구성되어 있습니다.

• step: 해당 parameter가 업데이트된 횟수를 나타냅니다. 여기서는 tensor([584.], device='cuda:2')로 되어 있는데, 이는 해당 parameter가 총 584번 업데이트 되었다는 의미입니다. (i.e., total training step = 584)