[Pytorh] AMP - Automatic Precision

AMP?

AMP는 Automatic Mixed Precision package의 약자로 모델의 single precision(FP32)를 두 종류의 precision(FP16, FP32)으로 학습하게 하여 빠르게 학습을 하게해주는 패키지이다.

모델의 Foward 연산은 서로 다른 두 행렬을 행렬 곱을 하는 것과 같다. 이 말은 각각의 두 행렬의 precision이 일치해야 동일한 값으로 간주하여 연산을 할 수 있다는 말이다. Input이 FP16이면 Operation을 FP16으로 진행하고(FP16으로 cast 가능한 연산에 한하여), Input이 FP32이면 Operation을 cast하지 않고 FP32으로 그대로 진행한다(또는 FP32로 연산할 수 없는 연산이면 그대로 FP32를 사용). Input이 FP16과 FP32로 혼용 되어 있는 상태이면, operation은 FP32로 연산하고 FP32 output 을 만든다.
⇒ FP16으로 연산할 수 있는 operation과 상황이 제한적으로 있으므로 mixed precision 이다.

model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

for input, target in data:
    optimizer.zero_grad()

    # Enables autocasting for the forward pass (model + loss)
    with autocast():
        output = model(input)
        loss = loss_fn(output, target)

    # Exits the context manager before backward()
    loss.backward()
    optimizer.step()

• Forward, Loss 계산은 Autocast with 단락에 두어, 연산을 빠르게 진행한다.
• Backward 연산은 Autocast 단락 밖에 두는 것을 권장하며, forward시의 precision을 그대로 따라간다. FP16 연산을 한 부분은 FP16을, FP32 연산을 사용한 부분은 FP32를 따라간다.
  loss.backaward() : loss를 기준으로 gradient를 계산.
  optimizer.step() : 위의 구한 gradient를 바탕으로 parameter update.

⇒ 위와 같이 Backward 연산이 Forward시에 사용한 precision을 사용하면 underflow 가 발생할 수 있는 문제 점이 있음.

GradScaler (torch.cuda.amp.GradScaler)

Forward pass를 FP16으로 하였다면, Backward pass 또한 FP16 gradient를 생성한다. Gradient의 값은 약간 크기 때문에 FP16으로 표현할 수 없을 수 있다. 이러한 현상을 underflow(0 이하의 값으로 넘친다)라고 표현한다. 이 현상을 해결하기 위해서 loss에 일정 크기의 scale factor를 곱해줌으로 써 backward pass시에 scaled loss로 연산을 하게하는 gradient scaling 을 한다.

Optimzier update를 하기 이전에 unscaled를 해주기 때문에, scale factor가 방해하는 요소로 작용하지 않는다.

scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
scaler.step(optimizer)
scaler.update()

• scaler.scale(loss) : loss에 내부적으로 존재하는 scale list를 각 loss에 맞는 값을 곱해준다.
• scaler.scale(loss).backward() : scaled loss 를 사용하여 backward pass를 진행하여 under flow가 발생하지 않는 gradient 계산.
• sclaer.unscale_(optimizer) : parameter update하기 이전에 gradient에 scale foctor 요소를 없애준다. 이 때 gradient가 inf, NaNs 체크도 같이 함.
• clip_grad_norm(..) :
• scaler.step(optimizer) : unscale_ 함수를 부르지 않았다면 내부적으로 호출하고, optimizer.step()을 unscaled gradient를 사용하여 업데이트 한다.
  ⇒ 따로 optimizer.step()을 호출하지 않아야 한다.
• scaler.update() : optimzier step 생략하더라도 backoff_factor를 호출하여 scale factor를 줄이고, growth_inerval가 반복적으로 발생한 경우 growth_factor를 곱하여 증가 시킨다.

참고 - Working ith Multiple Models, Lossesm and Optimizers

scaler = torch.cuda.amp.GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer0.zero_grad()
        optimizer1.zero_grad()
        with autocast(device_type='cuda', dtype=torch.float16):
            output0 = model0(input)
            output1 = model1(input)
            loss0 = loss_fn(2 * output0 + 3 * output1, target)
            loss1 = loss_fn(3 * output0 - 5 * output1, target)

        # (retain_graph here is unrelated to amp, it's present because in this
        # example, both backward() calls share some sections of graph.)
        scaler.scale(loss0).backward(retain_graph=True)
        scaler.scale(loss1).backward()

        # You can choose which optimizers receive explicit unscaling, if you
        # want to inspect or modify the gradients of the params they own.
        scaler.unscale_(optimizer0)

        scaler.step(optimizer0)
        scaler.step(optimizer1)

        scaler.update()