Gradient 누적

from transformers import TrainingArguments, Trainer
import torch
import gc

def gpu_memory_experiment(batch_size,
                        gradient_accumulation_steps=1,
                        gradient_checkpointing=False,
                        model_id="EleutherAI/polyglot-ko-1.3b",
                        peft=None):
    """
    GPU 메모리 사용량을 실험하는 함수
    
    Args:
        batch_size (int): 배치 크기
        gradient_accumulation_steps (int): 그래디언트 누적 스텝 수
        gradient_checkpointing (bool): 그래디언트 체크포인팅 사용 여부
        model_id (str): 허깅페이스 모델 ID
        peft (str): PEFT 방식 (예: 'qlora')
    
    Example:
        # 기본 학습
        gpu_memory_experiment(batch_size=4)
        
        # 그래디언트 누적 사용 (실제 배치 크기 = 4 * 4 = 16)
        gpu_memory_experiment(batch_size=4, gradient_accumulation_steps=4)
    """
    print(f"배치 사이즈: {batch_size}")
    
    # 모델과 토크나이저 로드
    model, tokenizer = load_model_and_tokenizer(model_id, peft=peft)
    
    # QLoRA나 그래디언트 체크포인팅 사용 시 캐시 비활성화
    if gradient_checkpointing == True or peft == 'qlora':
        model.config.use_cache = False
    
    # 더미 데이터셋 생성
    dataset = make_dummy_dataset()
    
    # 학습 인자 설정
    training_args = TrainingArguments(
        per_device_train_batch_size=batch_size,  # GPU당 배치 크기
        gradient_accumulation_steps=gradient_accumulation_steps,  # 그래디언트 누적 스텝 수
        gradient_checkpointing=gradient_checkpointing,  # 메모리 절약을 위한 체크포인팅
        output_dir="./result",
        num_train_epochs=1
    )
    
    try:
        # 모델 학습 시도
        train_model(model, dataset, training_args)
    except RuntimeError as e:
        # CUDA OOM 에러 처리
        if "CUDA out of memory" in str(e):
            print(e)
        else:
            raise e
    finally:
        # 메모리 정리
        del model, dataset
        gc.collect()
        torch.cuda.empty_cache()
        print_gpu_utilization()

"""
그래디언트 누적 작동 방식 예시:

batch_size=4, gradient_accumulation_steps=4인 경우:

1. 첫 번째 미니배치(4개 샘플)
   - 순전파 수행
   - 역전파로 그래디언트 계산
   - 그래디언트 저장 (가중치 업데이트 X)

2. 두 번째 미니배치(4개 샘플)
   - 순전파 수행
   - 역전파로 그래디언트 계산
   - 이전 그래디언트에 누적

3. 세 번째 미니배치(4개 샘플)
   - 순전파 수행
   - 역전파로 그래디언트 계산
   - 이전 그래디언트에 누적

4. 네 번째 미니배치(4개 샘플)
   - 순전파 수행
   - 역전파로 그래디언트 계산
   - 이전 그래디언트에 누적
   - 누적된 그래디언트로 가중치 업데이트

실제 효과:
- 메모리 사용량: batch_size=4 수준
- 계산 효과: batch_size=16과 동일
"""

1. 그래디언트 누적의 목적

   • 제한된 GPU 메모리로 더 큰 배치 크기 효과를 얻기 위해 사용
   • 작은 배치로 나눠서 계산하고 그래디언트를 모아서 한 번에 업데이트

2. 장점:

   • 메모리 효율적: 큰 배치를 직접 처리할 때보다 메모리 사용량이 적음
   • 큰 배치 효과: 작은 배치로 나눠 처리해도 큰 배치의 학습 효과를 얻을 수 있음
   • 학습 안정성: 큰 배치의 안정적인 그래디언트 효과를 얻을 수 있음

3. 실제 배치 크기 계산:

   • 효과적 배치 크기 = batch_size × gradient_accumulation_steps
   • 예: batch_size=4, gradient_accumulation_steps=4인 경우
     • 메모리는 4개 샘플만큼만 사용
     • 학습 효과는 16개 샘플 배치와 동일

4. 트레이드오프:

   • 학습 시간이 약간 증가할 수 있음
   • 메모리 사용량과 학습 속도 사이의 균형 필요

from torch.utils.data import DataLoader
from torch.optim import AdamW
import torch

def train_model(model, dataset, training_args):
    """
    모델 학습을 수행하는 함수
    
    Args:
        model: 학습할 모델
        dataset: 학습 데이터셋
        training_args: 학습 설정값들
        
    Example:
        model = AutoModelForCausalLM.from_pretrained("EleutherAI/polyglot-ko-1.3b")
        training_args = TrainingArguments(
            per_device_train_batch_size=4,
            gradient_accumulation_steps=4
        )
        train_model(model, dataset, training_args)
    """
    
    # 그래디언트 체크포인팅 활성화 (메모리 절약을 위해)
    if training_args.gradient_checkpointing:
        model.gradient_checkpointing_enable()
    
    # 데이터로더 설정
    train_dataloader = DataLoader(
        dataset, 
        batch_size=training_args.per_device_train_batch_size
    )
    
    # AdamW 옵티마이저 초기화
    optimizer = AdamW(model.parameters())
    
    # 학습 모드로 설정
    model.train()
    
    # GPU 사용량 한 번만 출력하기 위한 플래그
    gpu_utilization_printed = False
    
    # 학습 루프
    for step, batch in enumerate(train_dataloader, start=1):
        # 배치를 GPU로 이동
        batch = {k: v.to(model.device) for k, v in batch.items()}
        
        """
        예시) batch_size=4, gradient_accumulation_steps=4인 경우:
        
        Step 1:
        - inputs: {'input_ids': tensor([4개 샘플의 입력 ID들]), 'attention_mask': tensor([4개 샘플의 마스크])}
        - loss 계산 후 4로 나눔 (gradient_accumulation_steps)
        - backward() 호출하여 그래디언트 계산
        - 아직 optimizer.step() 호출하지 않음
        
        Step 2:
        - 새로운 4개 샘플로 위 과정 반복
        - 이전 step의 그래디언트에 누적
        
        Step 3, 4도 동일하게 진행
        
        Step 4 (마지막):
        - 4번째 미니배치 처리 완료
        - 누적된 그래디언트로 optimizer.step() 호출
        - optimizer.zero_grad()로 그래디언트 초기화
        """
        
        # 순전파 및 손실 계산
        outputs = model(**batch)
        loss = outputs.loss
        
        # 그래디언트 누적을 위해 loss를 gradient_accumulation_steps로 나눔
        loss = loss / training_args.gradient_accumulation_steps
        
        # 역전파로 그래디언트 계산
        loss.backward()
        
        # gradient_accumulation_steps만큼 그래디언트가 누적되면
        if step % training_args.gradient_accumulation_steps == 0:
            # 가중치 업데이트
            optimizer.step()
            
            # 메모리 사용량 추적
            gradients_memory = estimate_memory_of_gradients(model)
            optimizer_memory = estimate_memory_of_optimizer(optimizer)
            
            # GPU 사용량 출력 (첫 번째 업데이트 시에만)
            if not gpu_utilization_printed:
                print_gpu_utilization()
                gpu_utilization_printed = True
            
            # 그래디언트 초기화
            optimizer.zero_grad()
    
    # 최종 메모리 사용량 출력
    print(f"옵티마이저 상태의 메모리 사용량: {optimizer_memory / (1024 ** 3):.3f} GB")
    print(f"그레디언트 메모리 사용량: {gradients_memory / (1024 ** 3):.3f} GB")

"""
메모리 사용량 계산 예시:

1. 모델 파라미터 (가정):
   - 1.3B 모델의 경우 약 5.2GB의 파라미터
   - float32 타입: 4 bytes per parameter

2. 옵티마이저 상태:
   - AdamW는 각 파라미터에 대해 추가로 2개의 상태(m, v)를 저장
   - 메모리 사용량 ≈ 파라미터 크기 × 2
   - 1.3B 모델의 경우 약 10.4GB 추가 메모리 필요

3. 그래디언트:
   - 각 파라미터당 1개의 그래디언트 값
   - 메모리 사용량 ≈ 파라미터 크기
   - 1.3B 모델의 경우 약 5.2GB 추가 메모리 필요

4. 활성화(Activation) 메모리:
   - 배치 크기와 시퀀스 길이에 비례
   - gradient_checkpointing 사용 시 크게 감소
"""

세부 과정:

1. 그래디언트 누적의 세부 과정:

   • loss = loss / training_args.gradient_accumulation_steps:
     • 누적할 스텝 수만큼 손실을 나누어 그래디언트 스케일 조정
     • 이렇게 하지 않으면 누적된 그래디언트가 너무 커질 수 있음

2. 메모리 최적화 전략:

   • 그래디언트 체크포인팅: 활성화 메모리를 줄이는 대신 재계산
   • 배치 크기 조절: GPU 메모리와 학습 효율성의 균형
   • 그래디언트 누적: 효과적인 배치 크기 확보

3. 메모리 모니터링:

   • optimizer_memory: 옵티마이저의 상태 벡터들이 차지하는 메모리
   • gradients_memory: 모델 파라미터의 그래디언트가 차지하는 메모리

4. 최적화 팁:

   • gradient_accumulation_steps는 2의 배수로 설정하는 것이 일반적
   • 메모리 부족 시 gradient_checkpointing 활성화 고려
   • mixed precision training (fp16/bf16) 사용 고려