TPU 학습

TPU(Tensor Processing Unit)는 모델을 효율적으로 학습하도록 특별히 설계된 하드웨어 유형입니다. Accelerate는 TPU 학습을 지원하지만, 그래프 컴파일과 같이 알아야 할 몇 가지 사항이 있습니다. 이 튜토리얼에서는 컴파일을 간략히 다루며, 자세한 내용은 Accelerate를 사용한 TPU 학습 가이드를 참고하세요.

컴파일

TPU는 forward pass, backward pass, optimizer step과 같은 학습 단계의 모든 연산을 그래프로 생성합니다. 이 그래프를 빌드하고 컴파일하는 데 시간이 소요되기 때문에 첫 번째 학습 단계는 항상 시간이 걸립니다. 그러나 컴파일이 완료되면 결과가 캐시되어 이후의 모든 단계가 훨씬 빨라집니다.

핵심은 코드를 다시 컴파일하지 않는 것입니다. 그렇지 않으면 학습이 매우 느려집니다. 즉, 모든 연산이 정확히 동일해야 합니다:

• 배치의 모든 텐서는 동일한 길이를 가져야 합니다 (예: NLP 작업의 경우 동적 패딩을 사용하지 않음)
• your code는 정적이어야 합니다(입력에 따라 길이가 달라지는 for 루프가 있는 레이어, 예를 들어 LSTM과 같은 레이어는 사용할 수 없습니다).

Weight tying

일반적인 언어 모델 설계는 임베딩 레이어와 softmax 레이어의 가중치를 연결(tie)하는 것입니다. 그러나 모델을 TPU로 이동할 때(직접 이동하거나 prepare() 메서드에 전달하는 경우) 가중치 연결이 해제되므로 가중치를 다시 연결해야 합니다.

TPU를 위한 스크립트에 특수한 동작(가중치 연결 등)을 추가하려면 먼저 distributed_type을 DistributedType.TPU로 설정하세요. 그런 다음 tie_weights 메서드를 사용하여 가중치를 연결할 수 있습니다.

if accelerator.distributed_type == DistributedType.TPU:
    model.tie_weights()