테크니컬 스토리

      
        import os
        import torch
        import torch.distributed as dist
        import torch.multiprocessing as mp
        from torch.nn import Linear
        from torch.nn.parallel import DistributedDataParallel as DDP
        from torch.utils.data import DataLoader, TensorDataset, DistributedSampler

        #------------------

        #NCCL 환경 변수 설정

        #------------------

        os.environ["NCCL_DEBUG"] = "WARN"             # NCCL 경고 메시지 출력
        os.environ["NCCL_IB_DISABLE"] = "1"           # InfiniBand 비활성화
        os.environ["NCCL_P2P_LEVEL"] = "NVL"          # P2P 통신 레벨 설정
        os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"    # 사용할 GPU 선택

        #DDP 필수 환경 변수

        os.environ[\'MASTER_ADDR\'] = \'127.0.0.1\'
        os.environ[\'MASTER_PORT\'] = \'29500\'

        #------------------

        #DDP 초기화

        #------------------

        def setup(rank, world_size):
            dist.init_process_group("nccl", rank=rank, world_size=world_size)  # 프로세스 그룹 초기화
            torch.cuda.set_device(rank)   # 현재 프로세스를 특정 GPU에 매핑

        #------------------

        #DDP 종료

        #------------------

        def cleanup():
            dist.destroy_process_group()  # 프로세스 그룹 종료

        #------------------

        #학습 루프

        #------------------

        def train(rank, world_size):
            setup(rank, world_size)  # DDP 초기화

            # 모델 정의 및 GPU 이동
            model = Linear(10, 1).cuda(rank)  # Linear 모델을 해당 GPU로 이동
            ddp_model = DDP(model, device_ids=[rank])  # DDP로 래핑

            # 데이터셋 및 DataLoader
            dataset = TensorDataset(torch.randn(100, 10), torch.randn(100, 1))  # 랜덤 데이터 생성
            sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)  # DDP용 샘플러
            loader = DataLoader(dataset, batch_size=8, sampler=sampler)  # DataLoader 생성

            optimizer = torch.optim.SGD(ddp_model.parameters(), lr=0.01)  # 옵티마이저 설정

            # 학습 루프
            for epoch in range(2):

                running_loss = 0.0  # **추가: epoch별 누적 loss 초기화**
                for batch_idx, (x, y) in enumerate(loader):
                    x, y = x.cuda(rank), y.cuda(rank)  # 입력 데이터를 GPU로 이동
                    optimizer.zero_grad()              # 기울기 초기화
                    loss = ((ddp_model(x) - y) ** 2).mean()  # MSE Loss 계산
                    loss.backward()                    # 역전파
                    optimizer.step()                   # 가중치 업데이트
                    

                    running_loss += loss.item()  # **추가: 배치 loss 누적**

                    # 5배치마다 출력 (기존 출력문 수정)

                    print(f"[GPU {rank} | Epoch {epoch} | Batch {batch_idx}] Loss: {loss.item():.6f}")

                    # **기존: rank==0만 출력 → 이제 모든 GPU 출력**

                # ------------------- 추가: epoch 끝난 뒤 평균 loss 출력 -------------------
                avg_loss = running_loss / len(loader)
                print(f"[GPU {rank} | Epoch {epoch}] Average Loss: {avg_loss:.6f}")
                # **추가: 한 눈에 GPU별 학습 상태 확인 가능**

            cleanup()  # DDP 종료

        #------------------

        #메인 실행

        #------------------

        if __name__ == "__main__":
            world_size = torch.cuda.device_count()  # 사용 가능한 GPU 개수
            mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)  # 멀티프로세스 학습 시작
    

      
        #실행
        torchrun --standalone --nproc_per_node=2 ddp_test.py
    

      
        (pt-venv) test@ubuntu:~$ python ddp_test.py
        NCCL version 2.21.5+cuda12.4
        [GPU 1 | Epoch 0 | Batch 0] Loss: 0.512695
        [GPU 0 | Epoch 0 | Batch 0] Loss: 1.596358
        [GPU 0 | Epoch 0 | Batch 1] Loss: 3.234331
        [GPU 1 | Epoch 0 | Batch 1] Loss: 1.230204
        [GPU 0 | Epoch 0 | Batch 2] Loss: 1.251617
        [GPU 1 | Epoch 0 | Batch 2] Loss: 2.918502
        [GPU 0 | Epoch 0 | Batch 3] Loss: 0.790517
        [GPU 1 | Epoch 0 | Batch 3] Loss: 0.996800
        [GPU 0 | Epoch 0 | Batch 4] Loss: 1.685620
        [GPU 1 | Epoch 0 | Batch 4] Loss: 1.339857
        [GPU 0 | Epoch 0 | Batch 5] Loss: 0.529339
        [GPU 1 | Epoch 0 | Batch 5] Loss: 1.109558
        [GPU 0 | Epoch 0 | Batch 6] Loss: 0.962902
        [GPU 0 | Epoch 0] Average Loss: 1.435812
        [GPU 1 | Epoch 0 | Batch 6] Loss: 0.309195
        [GPU 1 | Epoch 0] Average Loss: 1.202401
        [GPU 1 | Epoch 1 | Batch 0] Loss: 0.476248
        [GPU 0 | Epoch 1 | Batch 0] Loss: 1.503265
        [GPU 1 | Epoch 1 | Batch 1] Loss: 1.115403
        [GPU 0 | Epoch 1 | Batch 1] Loss: 3.244143
        [GPU 1 | Epoch 1 | Batch 2] Loss: 2.682358
        [GPU 0 | Epoch 1 | Batch 2] Loss: 1.068382
        [GPU 1 | Epoch 1 | Batch 3] Loss: 0.868395
        [GPU 0 | Epoch 1 | Batch 3] Loss: 0.723319
        [GPU 1 | Epoch 1 | Batch 4] Loss: 1.248225
        [GPU 0 | Epoch 1 | Batch 4] Loss: 1.516317
        [GPU 1 | Epoch 1 | Batch 5] Loss: 0.936232
        [GPU 0 | Epoch 1 | Batch 5] Loss: 0.508594
        [GPU 1 | Epoch 1 | Batch 6] Loss: 0.269102
        [GPU 1 | Epoch 1] Average Loss: 1.085138
        [GPU 0 | Epoch 1 | Batch 6] Loss: 0.833279
        [GPU 0 | Epoch 1] Average Loss: 1.342471