분산 추론 개요¶

목차¶

• 분산이 필요한 이유
• 전략 비교
• 각 전략을 사용하는 시점
• 공통 인프라
• 클러스터 설정

분산이 필요한 이유¶

현대 LLM은 모델 가중치만으로도 수십에서 수백 기가바이트의 메모리가 필요합니다. 16비트 정밀도의 70B 파라미터 모델은 약 140 GB를 차지하므로 소비자용이든 워크스테이션용이든 GPU 한 장에는 올라가지 않습니다. 모델이 메모리에 맞더라도, 여러 요청을 동시에 처리하면 KV 캐시를 위한 장치 메모리가 고갈될 수 있습니다.

분산은 두 가지 별개의 문제를 해결합니다:

1. 메모리 용량: 모델이 하나의 GPU에 담기에 너무 큽니다. 해결책: 각 GPU가 일부분만 보유하도록 모델을 분할합니다.

2. 처리량: 단일 GPU가 원하는 요청 속도를 따라가지 못합니다. 해결책: 모델의 여러 복사본 또는 파이프라인 스테이지를 병렬로 실행합니다.

mlxcel은 상호 보완적인 세 가지 전략을 지원하며 배포 토폴로지에 따라 조합해 쓸 수 있습니다.

전략 비교¶

graph TB
    subgraph TP["Tensor Parallelism (TP)"]
        direction LR
        TP_IN[Input Tokens] --> TP_R0[Rank 0<br/>W shard 0]
        TP_IN --> TP_R1[Rank 1<br/>W shard 1]
        TP_R0 -->|all-reduce| TP_OUT[Output]
        TP_R1 -->|all-reduce| TP_OUT
    end

    subgraph PP["Pipeline Parallelism (PP)"]
        direction LR
        PP_S0[Stage 0<br/>Layers 0-15] -->|activations| PP_S1[Stage 1<br/>Layers 16-31]
        PP_S1 -->|activations| PP_S2[Stage 2<br/>Layers 32-47]
        PP_S2 --> PP_OUT[Logits]
    end

    subgraph DI["Disaggregated Inference (DI)"]
        direction LR
        DI_IN[Request] --> DI_PF[Prefill Node<br/>Compute KV cache]
        DI_PF -->|transfer KV| DI_DC[Decode Node<br/>Autoregressive decode]
        DI_DC --> DI_OUT[Tokens]
    end

각 전략을 사용하는 시점¶

텐서 병렬화를 사용하는 경우:

• 동일 호스트에 여러 GPU가 있거나 NVLink/Thunderbolt로 연결된 경우.
• 모델이 단일 GPU에 맞지 않지만 2-8개 GPU에 분산하면 맞는 경우.
• 요청당 낮은 지연이 필요한 경우 (TP all-reduce는 공유 메모리 버스에서 빠릅니다).
• 모델이 표준 트랜스포머, MoE, 또는 SSM 아키텍처인 경우 (모두 아키텍처 인식 샤드 플랜으로 지원됩니다).

파이프라인 병렬화를 사용하는 경우:

• 여러 호스트가 있고 모델이 어느 단일 호스트에도 맞지 않는 경우.
• 요청당 지연이 파이프라인 버블 오버헤드를 수용할 수 있는 경우.
• 요청이 예측 가능한 배치 크기로 도착하는 경우 (마이크로 배치가 파이프라인 버블을 효과적으로 상각합니다).
• 이기종 노드 집합에서 총 처리량을 극대화하려는 경우.

분리형 추론을 사용하는 경우:

• 프롬프트 길이가 크게 가변적인 워크로드를 처리하는 경우 (예: RAG, 긴 컨텍스트의 코드 완성).
• 프리필 노드와 디코드 노드의 하드웨어 프로파일이 다른 경우 (프리필에는 연산 집약적 GPU, 디코드에는 메모리 집약적 GPU).
• 다른 쪽을 변경하지 않고 프리필 용량과 디코드 용량을 독립적으로 확장하려는 경우.
• 긴 프롬프트가 이미 진행 중인 짧은 디코드 세션을 지연시키지 않아야 하는 경우.

공통 인프라¶

세 가지 전략 모두 src/distributed/에 정의된 공통 인프라 레이어를 공유합니다.

노드 검색¶

각 노드는 고유한 ID와 네트워크 주소로 자신을 알립니다. 초기 피어 목록은 클러스터 설정(TOML 또는 CLI 플래그)으로 지정합니다.

하트비트 및 장애 감지¶

백그라운드 하트비트 루프(src/distributed/heartbeat.rs)는 알려진 모든 피어에 주기적으로 핑을 전송합니다. 장애 감지기(src/distributed/failure_detector.rs)는 하트비트가 설정된 횟수만큼 누락되면 노드를 사용 불가로 표시하고 요청 라우터의 재라우팅을 트리거합니다.

연결 풀링¶

src/distributed/connection_pool.rs는 피어 노드에 대한 지속적인 TCP 연결을 유지합니다. 연결을 요청 간에 재사용하므로 중요 경로에서 TCP 핸드셰이크 오버헤드가 발생하지 않습니다.

전송 레이어¶

전송 추상화(src/distributed/transport.rs)는 다음을 지원합니다:

• TCP: 기본 전송 방식으로, 모든 네트워크 토폴로지에서 작동합니다.
• Thunderbolt: Thunderbolt로 연결된 Apple Silicon 장치를 위한 고대역폭 전송 방식입니다. src/distributed/thunderbolt_transport.rs를 참조하십시오.
• RDMA: 연결 시점에 피어별로 협상하는 선택적 가속 경로이며 TCP로 폴백합니다. src/distributed/rdma_transport.rs와 파이프라인 병렬화를 참조하십시오.

graph LR
    A[Node A] -->|TCP / Thunderbolt| B[Node B]
    B --> C[Node C]
    C --> A
    style A fill:#4a90d9
    style B fill:#4a90d9
    style C fill:#4a90d9

클러스터 시크릿, 핸드셰이크, 프레임 보호¶

파이프라인 역할과 disaggregated prefill, decode, router 역할은 사전 공유 클러스터 시크릿(--cluster-secret-file, MLXCEL_CLUSTER_SECRET_FILE, MLXCEL_CLUSTER_SECRET, 또는 클러스터 TOML의 [cluster] secret_file, 16바이트 이상) 없이는 시작하지 않습니다. 단일 노드 서빙(hybrid)은 시크릿을 읽지 않습니다. 시크릿은 세 곳에 쓰입니다.

• 핸드셰이크(src/distributed/cluster_handshake.rs, 프로토콜 버전 5). 모든 TCP 연결은 상호 challenge-response로 시작합니다. 양쪽이 전체 트랜스크립트에 대한 HMAC-SHA256으로 시크릿 보유를 증명하므로 잘못된 시크릿을 가진 피어는 어느 쪽에서든 거절됩니다. 수락하는 쪽은 클라이언트가 공지한 identity를 TCP 소스 IP에 묶습니다. 서버의 verdict(수락 또는 용량 초과)에는 자체 HMAC 태그가 있어서 경로상의 제3자가 한 verdict를 다른 verdict로 바꿀 수 없습니다.
• 프레임 보호(src/distributed/frame_protection.rs). 시크릿과 핸드셰이크 nonce에서 HKDF-SHA256으로 연결별 단방향 키를 유도합니다. 핸드셰이크 이후의 모든 프레임은 기본 suite인 AES-256-GCM으로 암호화되고 인증되며 암묵적 시퀀스 번호를 가지므로, 재전송, 순서 변경, 누락, 주입된 프레임은 인증에 실패하고 연결이 닫힙니다. 양쪽은 같은 suite를 써야 하며 핸드셰이크가 불일치를 거절합니다.
• 디스커버리 비콘(src/distributed/discovery_beacon.rs)도 같은 시크릿으로 서명합니다.

핸드셰이크를 마치지 않은 연결은 소스별 한도가 있는 별도의 pending 풀을 쓰므로(src/distributed/pending_handshakes.rs), 아무 데이터도 보내지 않는 소켓이 몰려도 이미 수립된 연결의 슬롯은 소진되지 않습니다.

텐서 프로토콜¶

src/distributed/tensor_protocol.rs는 노드 간 텐서 전송을 위한 와이어 형식을 정의합니다. 지원하는 데이터 타입:

분리형 추론에서 KV 캐시를 전송할 때는 선택적 LZ4 압축과 즉석 int8/int4 양자화로 대역폭 요구사항을 줄일 수 있습니다.

클러스터 설정¶

클러스터는 --cluster-config로 전달하는 TOML 파일로 설정하며 간단한 토폴로지라면 CLI 플래그(--node-id, --node-role, --peers)만으로도 설정할 수 있습니다.

TOML 설정¶

[cluster]
name = "my-cluster"
tensor_parallel_size = 2
pipeline_parallel_size = 1

[[nodes]]
id = "node-0"
address = "192.168.1.10:8080"
role = "tensor_parallel_rank"
rank = 0

[nodes.resources]
memory_bytes = 68719476736   # 64 GiB
compute_units = 0            # 0 = auto-detect

[[nodes]]
id = "node-1"
address = "192.168.1.11:8080"
role = "tensor_parallel_rank"
rank = 1

[nodes.resources]
memory_bytes = 68719476736
compute_units = 0

노드 역할¶

--node-role router는 추가로 disaggregated 배포의 모델 없는 HTTP 프런트를 선택합니다. TOML 노드 역할이 아니라 CLI 역할이며, 분리형 추론을 참조하십시오.

CLI 설정 (간단한 토폴로지)¶

# Node 0 (rank 0)
mlxcel-server --model ./models/llama-70b \
  --node-id node-0 \
  --node-role tensor_parallel_rank \
  --peers 192.168.1.11:8080

# Node 1 (rank 1)
mlxcel-server --model ./models/llama-70b \
  --node-id node-1 \
  --node-role tensor_parallel_rank \
  --peers 192.168.1.10:8080

검증¶

클러스터 설정은 시작 시 검증됩니다:

• 노드 ID는 고유하고 비어 있지 않아야 합니다.
• 주소는 고유해야 합니다.
• 노드 ID와 클러스터 이름은 128자를 초과하지 않아야 하며 제어 문자를 포함하면 안 됩니다.
• 최소 하나의 노드가 있어야 합니다.