[Transformers] IterableDatasetShard의 동작

 

분산학습 환경에서 IterableDataset을 사용하려고 하는데 데이터셋길이가 달라 문제가 발생한다.

보통의 경우에는 DataLoader에 drop_last=True 옵션을 주면 배치 사이즈가 다른 마지막 배치를 무시할 수 있다.

[ [[0, 1, 2], [3, 4, 5], [6, 7, 8]], [[9, 10, 11], [12, 13], [14, 15]] ]

 

그런데 LLM을 학습할 때 배치 사이즈가 1로 학습하면 아예 배치 자체의 길이가 달라지는 경우가 생긴다.

[ [[0], [1], [2]], [[3], [], [] ] <- 뒤의 두 배치가 아예 비어 버린다.

이렇게 뒤의 두 배치가 비어 버리면 첫번째 GPU는 [3]을 학습하는 step을 진행하고 있지만, 두번째, 세번째 GPU는 dataloader가 끝나면서 학습 loop를 종료하고 나간다.

이러면 첫번째 GPU는 loss = model(batch).loss 는 계산할 수 있지만, model.backward(loss)에서 다른 GPU의 gradient 계산을 기다리느라 멈춰버리는 것 같다.

 

대안은 transformers애서 제공하는 IterableDatasetShard를 쓰는 것이다.

4.35.2 버전을 기준으로 IterableDatasetShard는 항상 분산처리되는 배치의 길이를 동일하게 맞춰준다.

 

학습 데이터가 [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]이라고 했을 때, 3개 GPU에서 배치 사이즈가 2인 배치를 만들면

먼저 배치 사이즈 * GPU 개수만큼의 데이터를 뽑아내고, 그 데이터를 슬라이싱한다.

처음엔 [0, 1, 2, 3, 4, 5]가 뽑힌다.

그리고 이 현재 스텝의 총 배치 중 각각 [0, 1], [2, 3], [4, 5] 인덱스에 해당하는 데이터를 골라 보낸다.

그러면 [0, 1], [2, 3], [4, 5]가 각각 3개의 GPU에 할당된다.

 

그 다음엔 [6, 7, 8, 9]가 남아 있는데, 이는 필요한 총 배치 사이즈인 6보다 부족하다.

그래서 drop_last가 False인 경우, 현재 배치의 사이즈가 6보다 커질 때까지 자기 자신을 복제한다.

그러면, [6, 7, 8, 9, 6, 7, 8, 9]가 되고, 이 중에서 각각 [0, 1], [2, 3], [4, 5] 인덱스에 해당하는 데이터를 골라 보낸다.

그러면 [6, 7], [8, 9], [6, 7]이 각 GPU에 할당되어 모든 학습 프로세스가 잘 끝나게 된다.