PyTorch DDP non tensor data reduce

chatgpt로 작성한글입니다.

분산 환경에서 Non-Tensor 객체 동기화 방법

1. 배경

• 문제 상황:
  분산 학습 환경에서는 각 프로세스가 별도로 계산한 결과를 동기화할 필요가 있습니다.
  텐서(tensor)는 all_reduce나 all_gather 같은 연산을 통해 쉽게 동기화할 수 있지만, 파이썬의 일반 객체(예: 리스트, 딕셔너리 등)는 이러한 연산을 직접 사용할 수 없습니다.
• 해결 방법:
  PyTorch는 일반 파이썬 객체도 동기화할 수 있도록 torch.distributed.all_gather_object 함수를 제공합니다.

2. torch.distributed.all_gather_object 사용법

• 기능:
  각 프로세스에서 생성한 파이썬 객체를 한 곳으로 모으는 역할을 합니다.
• 핵심 포인트:
  • 각 프로세스가 동기화할 객체를 전달합니다.
  • 미리 결과를 저장할 리스트(수신 버퍼)를 준비해야 합니다.

3. 수신 버퍼 준비: [None for _ in range(dist.get_world_size())]

• 목적:
  모든 프로세스에서 보낸 객체들을 저장할 공간을 미리 확보하기 위함입니다.
• 설명:
  • dist.get_world_size()는 전체 프로세스의 수를 반환합니다.
  • 이 값을 이용해 전체 프로세스 수만큼의 요소를 가진 리스트를 만듭니다. 예를 들어: 
  • gathered_results = [None for _ in range(dist.get_world_size())]
  • 이렇게 하면 각 프로세스에서 전달한 객체가 해당 인덱스에 저장될 수 있습니다.

4. 전체 코드 예시

다음은 non-tensor 객체(예: 임의의 결과 리스트)를 모든 프로세스에서 동기화하는 예제입니다:

if distributed:
    # 각 프로세스의 결과를 담을 수신 버퍼 준비
    gathered_results = [None for _ in range(dist.get_world_size())]

    # 각 프로세스의 non-tensor 객체(예: 결과 리스트)를 동기화
    dist.all_gather_object(gathered_results, local_result)

    # 모든 프로세스에서 수집된 결과를 하나의 리스트로 합치기 (예시)
    combined_results = []
    for sublist in gathered_results:
        combined_results.extend(sublist)

• 과정 설명:
  1. 수신 버퍼 준비:
     [None for _ in range(dist.get_world_size())]를 사용하여 전체 프로세스 수만큼의 공간을 미리 할당합니다.
  2. 객체 동기화:
     dist.all_gather_object를 통해 각 프로세스의 객체(local_result)를 준비한 버퍼에 수집합니다.
  3. 결과 합치기:
     수집된 각 프로세스의 객체를 하나의 리스트(또는 다른 구조의 객체)로 합칩니다.

5. 결론

• 텐서가 아닌 일반 파이썬 객체는 all_reduce 등의 텐서 전용 동기화 메서드를 사용할 수 없습니다.
• 대신, torch.distributed.all_gather_object를 이용해 각 프로세스에서 계산한 non-tensor 객체를 안전하게 동기화할 수 있습니다.
• 이때, [None for _ in range(dist.get_world_size())]로 미리 수신 버퍼를 준비하는 것이 필수적입니다.