텐서 병렬 처리의 작동 방식

텐서 병렬 처리는 nn.Modules 수준에서 이루어집니다. 이는 모델의 특정 모듈을 텐서 병렬 랭크에 분할합니다. 이는 파이프라인 병렬 처리에 사용되는 모듈 세트의 기존 파티션에 추가됩니다.

텐서 병렬 처리를 통해 모듈을 분할하면 순방향 및 역방향 전파가 배포됩니다. 라이브러리는 이러한 모듈의 배포 실행을 구현하는 데 필요한 디바이스 간 통신을 처리합니다. 모듈은 여러 데이터 병렬 랭크로 분할됩니다. 기존 워크로드 배포와 달리 라이브러리의 텐서 병렬 처리가 사용되는 경우 각 데이터 병렬 랭크에는 전체 모델 복제본이 포함되지 않습니다. 대신, 각 데이터 병렬 랭크에는 배포되지 않은 전체 모듈 외에 배포 모듈의 파티션만 있을 수 있습니다.

예: 데이터 병렬 랭크에서의 텐서 병렬 처리를 생각해 봅니다. 여기에서 데이터 병렬도는 4, 텐서 병렬도는 2입니다. 모듈 세트를 분할한 후, 다음 모듈 트리를 포함하는 데이터 병렬 그룹이 있다고 가정해 봅시다.

A
├── B
|   ├── E
|   ├── F
├── C
└── D
    ├── G
    └── H

모듈 B, G, H에 텐서 병렬 처리가 지원된다고 가정합니다. 이 모델 텐서 병렬 분할의 가능한 결과 한 가지는 다음과 같습니다.

dp_rank 0 (tensor parallel rank 0): A, B:0, C, D, G:0, H
dp_rank 1 (tensor parallel rank 1): A, B:1, C, D, G:1, H
dp_rank 2 (tensor parallel rank 0): A, B:0, C, D, G:0, H
dp_rank 3 (tensor parallel rank 1): A, B:1, C, D, G:1, H

각 줄은 해당 dp_rank 안에 저장된 모듈 세트를 나타내고 표기법 X:y는 모듈 X의 y번째 부분을 나타냅니다. 유의할 사항:

1. 분할은 전체 DP_GROUP이 아닌 TP_GROUP이라고 하는 데이터 병렬 랭크의 서브셋에서 이루어지므로 정확한 모델 파티션이 dp_rank 0 및 dp_rank 2에 복제되고 dp_rank 1 및 dp_rank 3에도 동일하게 복제됩니다.

2. 모듈 E 및 F는 더 이상 모델의 일부가 아닙니다. 상위 모듈 B가 분할되어 있고 일반적으로 E 및 F의 일부인 실행은 (분할된) B 모듈에서 이루어지기 때문입니다.

3. H는 텐서 병렬 처리를 지원하지만 이 예제에서는 분할되지 않았습니다. 따라서 모듈을 분할할지 여부는 사용자 입력에 따라 달라집니다. 모듈이 텐서 병렬 처리를 지원한다고 해서 반드시 모듈이 분할된 것은 아닙니다.

라이브러리가 텐서 병렬 처리를 모듈에 적용하는 방법 PyTorch nn.Linear

데이터 병렬 랭크에 텐서 병렬 처리를 수행하면 파라미터, 그라디언트 및 옵티마이저 상태의 서브셋이 분할된 모듈의 텐서 병렬 디바이스 전체에서 분할됩니다. 나머지 모듈의 경우 텐서 병렬 디바이스가 일반적인 데이터 병렬 방식으로 작동합니다. 디바이스는 분할된 모듈을 실행하기 위해 먼저 동일한 텐서 병렬 처리 그룹에 있는 피어 디바이스의 모든 데이터 샘플에서 필요한 부분을 수집합니다. 그런 다음 디바이스는 이러한 모든 데이터 샘플에 대해 모듈의 로컬 부분을 실행한 다음, 또 다른 동기화 라운드를 실행합니다. 여기에서 각 데이터 샘플의 출력 부분을 결합하고 결합된 데이터 샘플을 데이터 샘플이 처음 생성된 GPU로 반환합니다. 다음 그림은 분할된 nn.Linear 모듈에 대한 이 프로세스의 예를 보여줍니다.

첫 번째 그림은 두 개의 텐서 병렬 처리 랭크에서 데이터 병렬 처리를 보유한 대형 nn.Linear 모듈이 있는 소형 모델을 보여줍니다. nn.Linear 모듈은 두 개의 병렬 랭크로 복제됩니다.

두 번째 그림은 nn.Linear 모듈을 분할하는 동안 더 큰 모델에 적용된 텐서 병렬 처리를 보여줍니다. 각 tp_rank 모듈에는 선형 모듈의 절반과 나머지 연산 전체가 포함됩니다. 선형 모듈이 실행되는 동안 각 tp_rank는 관련된 모든 데이터 샘플의 절반을 수집하여 nn.Linear 모듈의 절반을 통해 전달합니다. 각 랭크가 데이터 샘플의 최종 선형 출력을 보유하도록 결과를 감소-산란해야 합니다(합계를 감소 연산으로 사용). 나머지 모델은 일반적인 데이터 병렬 방식으로 실행됩니다.