파이프라인 병렬화¶

목차¶

• 개념
• 지원되는 family
• 제로 컨피그 시작
• 레이어 파티셔닝
• 스테이지 간 활성화 전송
• 마이크로 배치와 GPipe 스케줄
• 파이프라인 버블
• 스테이지별 KV 캐시 관리
• 트랜스포트 백엔드
• 설정
• 텐서 병렬화와의 조합 (2D PP × TP)
• PP와 LoRA

개념¶

파이프라인 병렬화(PP)는 모델을 레이어 기준으로 수직으로 분할합니다: 각 GPU(파이프라인 스테이지)는 연속된 범위의 트랜스포머 레이어를 담당합니다. 스테이지들은 파이프라인으로 배열됩니다. 스테이지 i는 배치 처리를 완료하고 결과 활성화를 스테이지 i+1에 전달하면 즉시 다음 마이크로 배치 처리를 시작할 수 있습니다.

이것은 모든 GPU가 모든 레이어에 참여하는 텐서 병렬화와 다릅니다. PP에서는 각 GPU가 자신의 스테이지에 할당된 레이어만 실행합니다.

graph LR
    subgraph "Pipeline (3 stages, 4 micro-batches)"
        subgraph S0["Stage 0 (layers 0-15)"]
            S0_MB0[MB 0]
            S0_MB1[MB 1]
            S0_MB2[MB 2]
            S0_MB3[MB 3]
        end

        subgraph S1["Stage 1 (layers 16-31)"]
            S1_MB0[MB 0]
            S1_MB1[MB 1]
            S1_MB2[MB 2]
            S1_MB3[MB 3]
        end

        subgraph S2["Stage 2 (layers 32-47)"]
            S2_MB0[MB 0]
            S2_MB1[MB 1]
            S2_MB2[MB 2]
            S2_MB3[MB 3]
        end

        S0_MB0 -->|activations| S1_MB0
        S0_MB1 -->|activations| S1_MB1
        S0_MB2 -->|activations| S1_MB2
        S0_MB3 -->|activations| S1_MB3

        S1_MB0 -->|activations| S2_MB0
        S1_MB1 -->|activations| S2_MB1
        S1_MB2 -->|activations| S2_MB2
        S1_MB3 -->|activations| S2_MB3
    end

스테이지들은 겹쳐서 동작합니다. 스테이지 1이 마이크로 배치 0을 처리하는 동안 스테이지 0은 이미 마이크로 배치 1을 처리하므로 노는 스테이지가 없습니다.

레포지토리는 현재 두 개의 PP 진입점을 제공합니다:

• mlxcel generate: 지원되는 텍스트 family(Llama, Mistral, Mixtral, DeepSeek V3, Llama 4 Scout 텍스트 타워, GPT-OSS, Gemma 3, Gemma 4, Qwen3, Qwen3.5, Qwen3-Next, GLM4 family, 그리고 하이브리드 Mamba + Transformer family인 Jamba와 Nemotron-H)에 대한 단일 머신 프로세스 내 파이프라인 실행.
• mlxcel serve: 코디네이터 주도 분산 서빙.

스테이지 로컬 실행은 다음과 같이 의도적으로 분리되어 있습니다:

• 토큰 임베딩, 레이어별 실행, 최종 norm, lm-head 방출을 위한 공유 스테이지 헬퍼
• 자기 family의 블록 타입만 로드하고 실행할 줄 아는 family별 백엔드

이 구조 덕분에 새 텍스트 family가 추가되어도 트랜스포트 프로토콜, 워커 루프, 원격 런타임은 변경되지 않습니다.

지원되는 family¶

각 지원 family에는 mlxcel-server가 클러스터 핸드셰이크 시점에 광고하는(src/distributed/pipeline/stage_executor/mod.rs의 supported_families()) 안정적인 텍스트 이름이 있습니다. 코디네이터와 스테이지 워커는 연결 시점에 이 이름들을 교환하며 워커가 코디네이터가 선언한 family를 포함하지 못하면 클러스터는 생성 결과를 조용히 손상시키는 대신 pipeline capability mismatch 오류를 내고 시작에 실패합니다.

대표 노트:

• Gemma 4는 다운스트림 per_layer_inputs를 스테이지 hidden-state 페이로드에 패킹해 트랜스포트와 런타임 프로토콜을 그대로 유지합니다.
• Gemma 4에는 KV 공유 레이어도 있으므로 대표 2-스테이지 분할은 각 공유 소스 레이어를 해당 소비자와 같은 스테이지에 보관해야 합니다. gemma-4-e2b-it-4bit의 경우 검증된 분할은 0-12,13-34입니다.
• 현재 GLM4 family 대표 검증은 glm4-flash-4bit를 사용하며 이는 GLM4 MoE Lite로 감지되어 전체 모델 로더가 사용하는 명시적 프리필 causal-mask 경로를 구동합니다.
• GLM MoE DSA는 공유된 DeepSeek V3.2 스테이지 모델 로더 경로를 재사용하지만 리포지토리는 현재 대표 실모델 패리티를 위한 로컬 glm5-4bit 체크포인트를 포함하지 않습니다.
• DeepSeek V3는 num_hidden_layers의 마지막 항목(MTP 트레일러)을 트랜스포머 블록 개수에서 제외합니다. 파이프라인 파티셔너는 num_hidden_layers - 1 레이어를 보므로 분할도 해당 범위 안에서 지정해야 합니다. sanitize_weights는 부분 로드 이전에 전체 가중치 맵에 대해 한 번 실행되므로 분할된 kv_b_proj 분해와 스택된 MoE expert 가중치가 모든 스테이지에서 일관되게 유지됩니다.
• Llama 4 Scout는 파이프라인 병렬화에서 텍스트 언어 타워만 노출합니다. 비전 인코더 스테이지 분할은 이번 롤아웃의 범위 밖입니다. 스테이지 실행기는 표준 KVCache를 사용하는 레거시 TransformerBlock::forward(x, cache, mask) 경로를 사용하여 트랜스포트 와이어 프로토콜을 그대로 유지합니다. 전체 모델의 iGQA ChunkedKVCache 최적화는 파이프라인 모드에서 사용되지 않습니다.
• Jamba와 Nemotron-H는 하이브리드 Mamba + Transformer 아키텍처입니다. 모든 Mamba / Mamba2 블록은 토큰마다 갱신되는 지속적인 conv_state + ssm_state를 유지하며 와이어 프로토콜이 hidden state만 전달하기 때문에 단일 SSM 블록을 여러 스테이지에 걸쳐 분할하는 것은 지원되지 않습니다. LayerFilter::layer_range는 블록 단위이므로 스테이지 경계는 두 디코더 블록 사이(어텐션 ↔ Mamba 포함)에는 놓일 수 있지만 블록 내부에는 놓일 수 없습니다. 두 하이브리드 실행기 모두 스테이지마다 전체 모델을 로드하고 로컬 스테이지에 할당된 레이어 범위만 실행합니다. 하이브리드 모델의 부분 가중치 로딩은 현재 구현 범위에 포함되지 않습니다. SSM 스테이지 경계의 정확성 불변식은 이 실행 방식과 무관합니다.
• Kimi K3(kimi_k3)는 지원되는 파이프라인 계열이 아닙니다. 93계층 전체 체크포인트는 4비트에서 약 1.4 TB로 어떤 단일 Apple Silicon 호스트보다 크지만 스테이지 실행기가 없으므로 --pp-size와 --pp-layers는 kimi_k3 모델을 거부합니다. 파티션 계획기는 dense, MLA, MoE 계층의 비용을 따로 계산합니다. 이 계열은 단일 프로세스로만 실행하며, 계층을 줄인 로컬 사본(config.json의 text_config.num_hidden_layers를 낮춘 것)이 로딩, 형상, 유한한 로짓을 단일 호스트에서 확인하는 방법입니다. 이 사본의 출력 텍스트에는 의미가 없습니다.

제로 컨피그 시작¶

제로 컨피그 구동은 CLI 플래그만으로 멀티 머신 파이프라인을 시작합니다. 클러스터 TOML을 손으로 쓸 필요도, 피어마다 --node-id를 줄 필요도 없습니다. 코디네이터가 클러스터를 계획하고 각 --pp-peer 프로세스에 스테이지 할당(스테이지 인덱스, 모델 식별 정보, 레이어 범위, 클러스터 토폴로지)을 푸시하며, 피어는 --distributed-config나 --node-id가 필요 없습니다. 플래그는 mlxcel-server와 대응되는 mlxcel serve 서브커맨드에서 쓸 수 있습니다:

코디네이터는 피어 주소를 문자열 사전 순으로 정렬한 뒤 파티셔닝하므로, 동일한 피어 집합으로 재실행하면 바이트 단위로 동일한 TOML이 생성됩니다. 이 파일은 재현성을 위해 버전 관리 시스템에 커밋해도 안전합니다. 내부 계획 기반 구현은 src/distributed/cluster_init.rs에 있으며 생성된 TOML은 기존 ServerStartupConfig 분산 처리 경로에서 그대로 소비됩니다. 즉 제로 컨피그 경로는 수동 TOML 경로와 동일한 런타임 코드를 공유하며 병렬 프로토타입이 아닙니다.

클러스터 시크릿과 전송 보호¶

모든 파이프라인 역할(--pp-auto, --pp-peer, pipeline_parallel_size >= 2인 --distributed-config)과 모든 disaggregated 역할은 클러스터 시크릿 없이는 시작하지 않습니다. 시크릿은 한 번 생성해 모든 호스트에 복사합니다.

openssl rand -hex 32 > cluster.secret

시크릿은 --cluster-secret-file, MLXCEL_CLUSTER_SECRET_FILE(경로), MLXCEL_CLUSTER_SECRET(값), 클러스터 TOML의 [cluster] secret_file 순서로 읽습니다. 시크릿은 세 곳에 쓰입니다. 모든 TCP 연결의 상호 HMAC-SHA256 핸드셰이크, 핸드셰이크 이후 모든 프레임을 AES-256-GCM으로 암호화하고 인증하는 연결별 키(HKDF-SHA256), 그리고 디스커버리 비콘 서명입니다.

핸드셰이크는 클러스터 프로토콜 버전 5입니다. v0.8.0보다 오래된 빌드는 v0.8.0 클러스터에 참여할 수 없으므로 모든 노드를 함께 업그레이드해야 합니다. 서버의 수락, 용량 초과, 거부 verdict도 인증되므로 경로상의 제3자가 한 verdict를 다른 verdict로 바꿀 수 없습니다. 프레임에는 암묵적 시퀀스 번호가 있어서 재전송, 순서 변경, 누락, 주입된 프레임은 인증에 실패하고 연결이 닫힙니다.

예시: 정적 시드를 사용한 LAN 구동¶

# 머신 B (피어): 192.168.1.11:19001에서 스테이지 할당을 기다립니다.
mlxcel-server --model ./models/llama-3.2-1b-4bit \
  --pp-peer \
  --cluster-secret-file ./cluster.secret \
  --cluster-control-addr 192.168.1.11:19001 \
  --host 0.0.0.0 --port 18081 --no-warmup

# 머신 C (피어): 192.168.1.12:19001에서 기다립니다.
mlxcel-server --model ./models/llama-3.2-1b-4bit \
  --pp-peer \
  --cluster-secret-file ./cluster.secret \
  --cluster-control-addr 192.168.1.12:19001 \
  --host 0.0.0.0 --port 18082 --no-warmup

# 머신 A (코디네이터): 클러스터를 계획하고 각 피어에 할당을 푸시합니다.
mlxcel-server --model ./models/llama-3.2-1b-4bit --alias llama-auto-pp \
  --pp-auto 2 \
  --cluster-secret-file ./cluster.secret \
  --cluster-peers 192.168.1.11:19001,192.168.1.12:19001 \
  --cluster-control-addr 192.168.1.10:19000 \
  --cluster-config-out ./.mlxcel/cluster.toml \
  --parallel 2 --max-batch-size 2 --pp-micro-batch-size 2 \
  --host 0.0.0.0 --port 18080

코디네이터는 스테이지를 직접 실행하지 않으므로 --pp-auto 2에는 --pp-peer 프로세스 2개가 필요합니다. 할당 교환은 세 단계입니다. 각 피어가 자신의 offer를 수락하거나 거절하고, 모든 피어가 수락한 뒤에야 코디네이터가 commit을 보내며, 모든 피어가 commit한 뒤에야 activate를 보냅니다. 피어는 activate를 받으면 클러스터 TOML을 쓰고 스테이지를 시작합니다. 한 피어라도 거절하면 어디서도 아무것도 시작하지 않으며, 체크포인트가 offer의 모델 정보(config 해시, 가중치 파일 이름과 크기, 레이어 수)와 맞지 않는 피어는 타입이 지정된 오류로 거절합니다. 코디네이터는 모든 피어가 commit한 뒤에 HTTP 서빙을 시작합니다.

예시: 옵트인 mDNS 디스커버리¶

피어를 사전에 열거할 수 없을 때는, 코디네이터가 피어의 UDP 비콘을 기다리도록 할 수 있습니다:

# 각 피어는 기본 비콘 포트에서 서명된 UDP 비콘으로 자신을 알립니다.
mlxcel-server --model ./models/llama-3.2-1b-4bit \
  --pp-peer \
  --cluster-secret-file ./cluster.secret \
  --cluster-discovery=mdns \
  --cluster-name studio-pp

# 코디네이터는 피어 2개가 공지할 때까지 기다린 뒤 토폴로지를 계획합니다.
mlxcel-server --model ./models/llama-3.2-1b-4bit \
  --pp-auto 2 \
  --cluster-secret-file ./cluster.secret \
  --cluster-discovery=mdns \
  --cluster-name studio-pp \
  --cluster-config-out ./.mlxcel/cluster.toml

mDNS 디스커버리는 기본적으로 비활성화되어 있습니다. 활성화하면 코디네이터는 자신의 --cluster-name에 대해 클러스터 시크릿으로 서명된 비콘만 받아들이므로, 같은 LAN에 있는 다른 애플리케이션의 브로드캐스트 트래픽은 무시됩니다.

Dry run과 오퍼레이터 오류¶

--dry-run은 해결된 토폴로지 요약과 방출된 TOML의 경로를 출력하되, 모델 로드와 HTTP 리스너를 시작하지 않습니다:

$ mlxcel-server -m ./models/llama-3.2-1b-4bit \
    --pp-auto 2 \
    --cluster-peers 192.168.1.11:19001,192.168.1.12:19001 \
    --dry-run
Cluster: mlxcel-cluster
  TP size: 1, PP size: 2
  Nodes (3):
    - coordinator @ 127.0.0.1:19000 [hybrid]
    - stage-0 @ 192.168.1.11:19001 [pipeline_stage]
    - stage-1 @ 192.168.1.12:19001 [pipeline_stage]
Emitted cluster TOML at: /home/op/.mlxcel/cluster.toml

Dry run은 다음 상황에서 0이 아닌 종료 코드로 종료됩니다:

• --pp-auto N이 설정되었지만 (정적 또는 mDNS 경로에서) 해결된 피어 수가 N보다 적을 때.
• 코디네이터 제어 주소가 HTTP 리슨 주소와 충돌할 때.
• 계획된 토폴로지의 TCP 포트가 이미 다른 프로세스에 의해 바인딩되어 있을 때.

오류 메시지에는 실행 가능한 해결 힌트(어떤 플래그를 바꿀지, 어떤 서브넷/방화벽 규칙을 확인할지, 어떤 폴백을 사용할지)가 포함됩니다.

결정성 보장¶

클러스터 TOML은 동일 토폴로지에서 바이트 단위로 동일함이 보장됩니다. 구체적으로:

• 피어는 스테이지 할당 전에 문자열 주소를 기준으로 사전 순으로 정렬됩니다. 따라서 --cluster-peers A,B와 --cluster-peers B,A는 동일한 TOML을 생성합니다.
• 노드 키는 고정된 순서(id, address, role, stage, rank)로 방출되며 코디네이터 항목은 항상 파이프라인 스테이지보다 앞에 나타납니다.
• 헤더 주석은 고정된 문자열이며 코디네이터의 호스트명, PID, 시계 값을 포함하지 않습니다.

결정성 보장은 src/distributed/cluster_init_tests.rs의 plan_cluster_emits_byte_identical_toml, plan_cluster_reordered_peers_produce_same_topology, render_deterministic_toml_is_stable 테스트로 검증합니다.

수동 TOML 경로에서 마이그레이션¶

기존의 수동 TOML / scripts/benchmark_pipeline_remote_rollout.sh write-config 워크플로는 제거되지 않습니다. 제로 컨피그 플래그는 가산적으로 조합됩니다. 이미 동작하는 클러스터 TOML이 있는 오퍼레이터는 계속 --distributed-config를 사용하고 --pp-auto를 무시해도 됩니다. 두 경로를 한 번의 호출에서 혼용하면 설명이 포함된 오류와 함께 시작 시 거부됩니다.

레이어 파티셔닝¶

레이어 파티셔닝은 src/distributed/pipeline/partition.rs(핵심 타입과 검증기), partition_balance.rs(바이트 정확 제약 조건 기반 밸런서), partition_profile.rs(모델 인식 ModelProfile 빌더), partition_quality.rs(추정치 대 실측 메모리 리포트)에 구현되어 있습니다.

자동 파티셔닝 (기본값)¶

자동 파티셔너는 스테이지별 최대 바이트 적재량을 최소화하도록 분배하되, 네 가지 강제 제약 조건을 준수합니다. 밸런서는 바이트 정확 + 제약 조건 인식 방식이어서 MoE 모델과 Gemma 4 / Gemma 4 VLM에서도 운영자가 --pp-layers를 수동으로 지정하지 않아도 기본 경로가 동작합니다.

알고리즘:

1. config.json을 읽어 실제 레이어별 바이트 가중치와 인접성 제약을 담은 ModelProfile을 구축합니다. MoE 레이어는 전문가 블록을 포함하여 계산하고 Gemma 4의 KV 공유 consumer 레이어는 더블 와이드 MLP를 반영 하며 양자화 비트 폭으로 추정치를 스케일링합니다.
2. 스테이지 0에 임베딩, 마지막 스테이지에 lm_head를 예약합니다. 나머지 스테이지는 레이어 전용 바이트 예산만 받습니다.
3. 모든 스테이지별 예산을 지키고 LayerAdjacencyGroup을 잘라내지 않는 K 스테이지 분할이 존재하는 최소 T(스테이지별 최대 바이트)를 이진 탐색합니다. 예산이 타이트한 스테이지는 더 적은 바이트를 받고 남은 가중치는 여유 있는 스테이지로 흘러갑니다.
4. 가능한 T에 대해 레이어 범위를 재구성합니다. 스테이지 0에 임베딩, 마지막 스테이지에 lm_head를 할당하며 활성화 전송 비용을 줄이기 위해 레이어는 연속된 범위로 유지합니다.

모델 인식 레이어별 바이트 계산¶

모든 레이어에 동일한 layer_param_bytes를 적용하면 MoE 계열에서 편향이 생깁니다. 전문가가 많은 레이어를 옆의 dense 레이어와 동일하게 계산하므로 밸런서가 레이어를 엉뚱한 스테이지에 배정하고, 무거운 스테이지는 로드 시 메모리 부족을 겪습니다. 프로필 빌더는 대신 다음을 인식합니다:

• Mixtral / Qwen 2.5 MoE / GLM 4 MoE / GLM 4 MoE Lite / GLM MoE DSA / GPT-OSS / Phi-MoE / Exaone MoE: 모든 레이어가 MoE이며 레이어별 바이트에 num_local_experts 전문가 가중치 텐서가 모두 포함됩니다.
• DeepSeek V3 / Nemotron-H: first_k_dense_replace가 dense 레이어 범위를, moe_layer_freq가 MoE 레이어 주기를 결정합니다. 프로필은 각 레이어를 그에 맞춰 표시합니다.
• Llama 4 (텍스트 타워 / VLM): interleave_moe_layer_step마다 intermediate_size 전문가를 가진 레이어를 MoE로 선택하고 dense 레이어는 intermediate_size_mlp를 사용합니다.
• Jamba: expert_layer_period / expert_layer_offset가 MoE 레이어를 표시합니다. Mamba 블록 자체는 스테이지 로컬로 유지됩니다.
• Gemma 4 / Gemma 4 VLM: num_kv_shared_layers + layer_types에서 KV 공유 인접성 그룹을 생성하며(아래 참조), use_double_wide_mlp가 consumer 레이어의 MLP 풋프린트를 두 배로 늘립니다.
• 양자화 모델: quantization.bits(4비트, 8비트 등)를 준수하며 스케일/바이어스에 대한 약간의 오버헤드를 더합니다.

KV 공유 레이어 인접성 (Gemma 4)¶

Gemma 4의 마지막 num_kv_shared_layers 디코더 블록은 자체 key/value를 계산하지 않고 layer_types[i] 값이 동일한 이전 레이어의 캐시를 빌려 씁니다. source 레이어와 consumer 사이를 분할하면 캐시가 엉뚱한 스테이지에 남아 생성이 깨집니다.

프로필 빌더는 모든 (source, consumer) 쌍을 LayerAdjacencyGroup에 기록하며 밸런서는 해당 그룹 내부에 스테이지 경계가 떨어지는 분할을 거부합니다. 그 결과 mlxcel generate --pp-size 2와 mlxcel-server --pp-auto 2는 gemma-4-e2b-it-4bit에 대해서도 기본값으로 유효한 플랜을 만듭니다.

수동 파티셔닝 (계속 지원)¶

레이어 범위는 --pp-layers로 명시적으로 지정할 수 있습니다:

# 3-stage pipeline; stage 0 gets layers 0-19, stage 1 gets 20-39, stage 2 gets 40-59
mlxcel-server --model ./models/llama-60b \
  --pp-layers "0-19,20-39,40-59" \
  ...

같은 문법은 CLI generate 경로에서도 사용할 수 있습니다:

mlxcel generate \
  --model ./models/llama-3.2-1b-4bit \
  --prompt "Hello" \
  --max-tokens 64 \
  --temp 0 \
  --no-chat-template \
  --pp-layers "0-7,8-15"

언제 수동 오버라이드가 유용합니까?

• 재현 가능한 실험: 프로필 빌더가 바뀌어도 실행 간에 바이트 단위로 동일한 할당이 필요한 경우.
• 이기종 플릿: 메모리 추정과 무관하게 특정 스테이지가 특정 범위를 가져야 하는 경우.
• 디버깅: 특정 경계가 버그를 재현하거나 고치는지 검증하려는 경우.

수동 플랜도 인접성 검증을 거치며 Gemma 4 KV 공유 그룹을 잘라내려고 하면 어떤 source / consumer 쌍이 깨지는지 가리키는 명시적 오류가 발생합니다.

파티션 품질 리포트¶

파티셔너는 스테이지별 추정 바이트 적재량, 전체 불균형 비율(max / min), 인접성 제약을 지키는 동안 발생한 경고를 모은 PartitionQualityReport를 생성합니다. 서버 모드에서 --verbose로 상세 로깅을 켜면 mlxcel-server가 시작 시 한 번 여러 줄짜리 블록을 로깅합니다:

pipeline partition quality: 2 stages, max=12500000 bytes, min=11900000 bytes, imbalance=1.5x
  stage 0 device 'local-stage-0' layers=13 estimated=12500000 actual=unknown
  stage 1 device 'local-stage-1' layers=13 estimated=11900000 actual=unknown

actual 열은 각 스테이지 가중치가 로드되면 채워지므로, 운영자는 바이트 비용 추정기가 실제와 얼마나 어긋나는지 확인할 수 있습니다(drift=... 접미사). 또한 밸런서는 인접성 제약이 크게 불균형한 플랜을 강제하는 경우 (max/min이 약 1.5배를 넘을 때) 파티션 시점에 명시적 경고를 발생시킵니다:

pipeline auto-partition produced an imbalanced plan: stage byte loads range from
80 to 120 (1.50x), which may hurt throughput. Heaviest stage 1 carries
adjacency group 0..7 (synthetic forced-join).

이 경고는 하드웨어 재조정을 고려하라는 힌트이며 오류가 아닙니다. 생성은 계속 진행됩니다.

스테이지 할당¶

각 StageAssignment는 다음을 기록합니다:

• stage_index: 0 기반 스테이지 번호.
• device_id: 클러스터 설정의 노드 ID와 일치.
• layer_range: 레이어 인덱스의 반열린 범위 [start, end).
• has_embedding: 스테이지 0의 경우 true.
• has_lm_head: 마지막 스테이지의 경우 true.
• estimated_memory_bytes: 이 스테이지의 추정 총 메모리.

스테이지 간 활성화 전송¶

스테이지가 마이크로 배치 처리를 완료하면, 전송 레이어를 거쳐 출력 활성화를 다음 스테이지로 전송합니다.

활성화 텐서는 TensorProtocol 와이어 형식(src/distributed/tensor_protocol.rs)으로 직렬화되어 TCP 또는 Thunderbolt 연결로 전송됩니다. 전송은 보내는 쪽 스테이지의 다음 마이크로 배치 순전파와 비동기적으로 진행됩니다.

sequenceDiagram
    participant S0 as Stage 0
    participant NET as Network
    participant S1 as Stage 1

    S0->>S0: Forward MB 0 (layers 0-15)
    S0->>NET: Send activations MB 0
    S0->>S0: Forward MB 1 (layers 0-15)
    NET->>S1: Receive activations MB 0
    S1->>S1: Forward MB 0 (layers 16-31)
    S0->>NET: Send activations MB 1
    S1->>S1: Forward MB 1 (layers 16-31)

활성화 전송 모듈은 src/distributed/pipeline/activation_transfer.rs에 있습니다.

마이크로 배치와 GPipe 스케줄¶

전체 요청 배치는 파이프라인에 진입하기 전에 마이크로 배치로 분할됩니다. 이것이 모든 스테이지를 동시에 바쁘게 유지하는 핵심 메커니즘입니다.

gantt
    title GPipe Pipeline Schedule (3 stages, 4 micro-batches)
    dateFormat  s
    axisFormat  %S

    section Stage 0
    MB 0 :a0, 0, 1s
    MB 1 :a1, after a0, 1s
    MB 2 :a2, after a1, 1s
    MB 3 :a3, after a2, 1s

    section Stage 1
    bubble :crit, b_bubble0, after a0, 1s
    MB 0 :b0, after b_bubble0, 1s
    MB 1 :b1, after b0, 1s
    MB 2 :b2, after b1, 1s
    MB 3 :b3, after b2, 1s

    section Stage 2
    bubble :crit, c_bubble0, after a0, 2s
    MB 0 :c0, after c_bubble0, 1s
    MB 1 :c1, after c0, 1s
    MB 2 :c2, after c1, 1s
    MB 3 :c3, after c2, 1s

GPipeSchedule은 실행 루프를 구동하는 ScheduleAction 값을 발생시킵니다:

마이크로 배치 분할은 src/distributed/pipeline/micro_batch.rs에 있습니다.

권장 마이크로 배치 크기¶

suggested_micro_batch_size(batch_size, num_stages)는 파이프라인 활용률을 극대화하기 위한 마이크로 배치당 권장 시퀀스 수를 반환합니다:

• num_stages 스테이지와 batch_size 시퀀스의 총 배치에서, 최적 마이크로 배치 수는 num_stages입니다 (모든 버블 슬롯을 채우기 위해).
• 각 마이크로 배치에는 ceil(batch_size / num_stages) 개의 시퀀스가 있습니다.

파이프라인 버블¶

파이프라인 버블은 스테이지가 이전 스테이지의 활성화를 기다리느라 처리할 작업이 없는 시간 슬롯입니다.

마이크로 배치가 m개인 k-스테이지 파이프라인에서 버블 비율은 다음과 같습니다:

버블을 최소화하려면: - 더 많은 마이크로 배치를 사용하십시오 (더 높은 m). m >> k이면 버블 비율이 0에 수렴합니다. - 더 작은 마이크로 배치를 사용하여 파이프라인을 더 빠르게 채우십시오.

예를 들어 마이크로 배치가 16개인 4-스테이지 파이프라인의 버블 비율은 3/18 ≈ 16.7%이고 32개면 3/34 ≈ 8.8%로 줄어듭니다.

graph LR
    subgraph "Bubble fill-in with more micro-batches"
        B[More micro-batches] --> LB[Lower bubble fraction]
        B --> HT[Higher throughput]
        LB --> HT
    end

스테이지별 KV 캐시 관리¶

각 파이프라인 스테이지는 자신이 소유한 레이어에 대한 KV 캐시를 유지합니다.

캐시 관리자(src/distributed/pipeline/cache_manager.rs)는 다음을 담당합니다:

• 스테이지당 시퀀스별 KV 캐시 항목을 할당하고 해제합니다.
• OOM을 방지하기 위해 스테이지당 캐시 용량을 추적합니다.
• 시퀀스가 선점될 때 캐시 축출를 조율합니다.

각 스테이지는 자신의 레이어 범위에 대한 캐시만 보유하므로, 총 KV 캐시 메모리가 단일 노드에 집중되지 않고 스테이지에 걸쳐 분산됩니다. 이것은 장문 컨텍스트 워크로드에서 파이프라인 병렬화의 메모리 이점 중 하나입니다.

트랜스포트 백엔드¶

스테이지 간 트랜스포트는 클러스터 TOML의 [cluster] 섹션에 있는 transport_backend 필드 또는 그에 대응하는 CLI 플래그로 선택합니다. 세 가지 백엔드가 기본으로 제공되며 모두 동일한 Transport 트레이트를 구현하므로 src/distributed/pipeline/activation_transfer.rs의 활성화 전송 소비자는 어떤 백엔드를 사용하더라도 변경되지 않습니다.

선택 기준¶

예시:

[cluster]
name = "pp-cluster"
pipeline_parallel_size = 2
transport_backend = "rdma"

[[nodes]]
id = "stage-0"
address = "192.168.1.11:9000"
role = "pipeline_stage"
stage = 0

[[nodes]]
id = "stage-1"
address = "192.168.1.12:9000"
role = "pipeline_stage"
stage = 1

RDMA 인식 백엔드¶

RDMA 인식 백엔드는 src/distributed/rdma_transport.rs에 구현되어 있으며 기존 TCP 코어 위에 얹혀 동작합니다. bind() 시점에 OS가 제로카피 프리미티브를 제공하는지 탐지(probing)합니다:

• Linux: io_uring 등록 버퍼. 탐지기는 먼저 커널 6.6+에 도입된 /proc/sys/kernel/io_uring_disabled sysctl을 조회한 뒤, 가벼운 io_uring_setup 시스템 호출을 시도하여 시스템 호출 자체에 접근할 수 있는지, seccomp / LSM이 차단하지 않는지 확인합니다.
• macOS: 등록된 메모리 영역과 kqueue + writev 배치 전송. Thunderbolt Bridge 인터페이스에 라우팅 가능한 IPv4 주소가 있으면 래퍼는 기존 Thunderbolt 트랜스포트 위에 붙어 검증된 인터페이스에서 등록 버퍼 전송을 수행합니다. 그렇지 않으면 일반 TCP 위에 붙습니다.
• 기타 플랫폼: 가속 경로를 시도하지 않습니다. os=<family>, driver=unsupported 사유와 함께 TCP로 폴백합니다.

피어들은 connect() 시점에 프로토콜 버전 핸드셰이크(RDMA_PROTOCOL_VERSION)를 교환합니다. 피어가 0(RDMA 인식 아님) 또는 불일치 버전을 광고하면 로컬 쪽은 피어 단위 폴백을 기록하고 TCP 시맨틱으로 계속 동작합니다.

조용하지만 기록되는 폴백¶

모든 폴백은 한 줄로 구조화된 사유(reason)와 함께 정확히 한 번 기록됩니다. 운영자는 로그 스트림에서 rdma_transport: falling back to TCP 접두사를 검색하면 되며 각 줄에는 OS / 드라이버 / 피어 버전 / 기능 불일치 형태의 reason= 키가 포함됩니다:

rdma_transport: falling back to TCP (reason=os=linux, driver=io_uring_missing: kernel does not export io_uring_setup (ENOSYS))
rdma_transport: peer negotiation fallback (peer=10.0.0.1:9000, reason=os=linux, peer_version=99: incompatible with local RDMA protocol version 1)

폴백이 발생해도 코디네이터는 계속 동작합니다. mlxcel generate / mlxcel-server는 제로카피 프리미티브 부재로 절대 크래시하지 않습니다. RdmaTransport::acceleration(), fallback_reason(), peers_negotiated(), peer_fallbacks()로 노출되는 구조화된 메트릭을 사용하면 협상 상태를 분산 메트릭 수집기로 쉽게 배관할 수 있습니다.

벤치마킹¶

2-노드 루프백 활성화 전송 벤치마크는 src/distributed/bench_activation.rs에 있고 tests/rdma_transport_bench.rs에서 구동됩니다. 실행 방법:

cargo test --release --test rdma_transport_bench \
    rdma_vs_tcp_activation_transfer -- --ignored --nocapture

하네스는 실제 활성화 형태를 반영한 네 가지 페이로드 크기(RPC 수준 핸드오프, 디코드 스텝 활성화, 프리필 크기 활성화, 스테이지 핸드오프 최악 케이스)에 대해 TCP와 RDMA의 단방향 지연(p50/p95/p99)과 대역폭을 마크다운 표로 출력합니다. 전송 경로를 변경할 때는 이 표를 로컬 기준값으로 사용하십시오.

설정¶

PipelineConfig는 최소 2개의 스테이지와 마이크로 배치당 최소 1개의 시퀀스를 요구합니다. 이 제약이 충족되지 않으면 시작 시 오류가 반환됩니다.

예시: 제로 컨피그 경로의 4-스테이지 파이프라인. --pp-peer 머신 4대가 필요하고 코디네이터는 스테이지를 실행하지 않습니다.

# 피어 4대 각각 (머신당 한 명령, --cluster-control-addr에는 자기 주소).
mlxcel-server --model ./models/llama-70b \
  --pp-peer \
  --cluster-secret-file ./cluster.secret \
  --cluster-control-addr 192.168.1.11:19001

# 코디네이터: 클러스터를 계획하고 각 피어에 스테이지 할당을 푸시합니다.
mlxcel-server --model ./models/llama-70b \
  --pp-auto 4 \
  --cluster-secret-file ./cluster.secret \
  --cluster-peers 192.168.1.11:19001,192.168.1.12:19001,192.168.1.13:19001,192.168.1.14:19001 \
  --cluster-config-out ./.mlxcel/cluster.toml

수동 다중 명령 워크플로는 생성된 클러스터 TOML과 표준 분산 플래그를 사용합니다:

# 공유 클러스터 TOML을 사용하는 stage 노드
mlxcel-server --model ./models/llama-70b \
  --distributed-config /shared/.mlxcel/cluster.toml \
  --node-id stage-0 \
  --node-role pipeline_stage \
  --peers 192.168.1.11:8080,192.168.1.12:8080,192.168.1.13:8080

# 다른 노드에서는 각 stage의 node-id로 반복합니다

텐서 병렬화와의 조합 (2D PP × TP)¶

파이프라인 병렬화는 텐서 병렬화와 조합해 2D 메시를 구성할 수 있습니다. 각 노드는 (pp_stage, tp_rank)의 한 교차점을 소유합니다. PP가 모델을 스테이지 단위로 수직 분할하고 한 스테이지 내에서 TP가 해당 스테이지의 레이어 가중치를 랭크에 걸쳐 수평으로 분할합니다.

참고: 동일한 메시를 TP 관점에서 서술한 내용은 텐서 병렬화 → 파이프라인 병렬화와의 조합 (2D PP × TP) 에서 확인할 수 있습니다.

언제 2D 병렬화를 사용하는가¶

PP-only 토폴로지에서 한 스테이지가 여전히 한 노드의 메모리를 초과할 때 사용합니다. 대표 사례는 Llama-70B 급 모델을 중형 GPU 4대 위에 2 PP 스테이지 × 2 TP 랭크 메시로 서빙하는 것입니다. PP가 모델을 두 개의 절반으로 나누고 TP가 각 절반을 다시 두 디바이스에 샤딩하므로 각 디바이스는 전체 가중치의 ¼만 보유합니다.

현재 제약:

• 스테이지별 TP 폭은 균일이어야 합니다. 이종 폭(예: 스테이지 0이 tp=4, 스테이지 1이 tp=2)은 지원하지 않습니다.
• 3D 병렬화는 범위 밖입니다. Expert 병렬화(PP × TP × EP)는 지원하지 않습니다.

역할 인코딩¶

2D 메시에 참여하는 노드는 pipeline_tensor_parallel 역할(pp_tp, 2d 별칭)을 사용합니다. 해당 노드는 TOML 엔트리에 stage와 rank를 모두 설정해야 합니다. 다음 조건이 검증 시 강제됩니다:

• 정확히 pipeline_parallel_size × tensor_parallel_size 개의 PPTP 노드
• 메시의 (stage, rank) 쌍이 정확히 한 번씩 존재
• PPTP가 아닌 노드는 stage/rank를 지정할 수 없음

역할 인코딩은 Rust에서 NodeRole::PipelineTensorParallel 열거형 변형과 캐시 측의 PpTpCoord 좌표로 노출됩니다.

트랜스포트 트래픽 클래스¶

2D 라우팅 레이어는 메시의 모든 노드에 걸쳐 단일 연결 풀을 공유하지만 메시지마다 올바른 핸들러가 처리하도록 두 트래픽 클래스를 구분합니다:

두 클래스 모두 PipelineTensorParallelRouter를 통해 선택됩니다. 이 라우터는 (preferred_stage, preferred_rank, traffic_class) 로 라우팅하여 단일 (stage, rank) 피어를 반환합니다. 한 스테이지의 모든 랭크(all-reduce용)나 같은 랭크의 모든 스테이지(PP 활성화 파이프라인용)로 브로드캐스트할 때는 NodeRegistry::nodes_at_stage 또는 nodes_at_rank를 직접 사용합니다.

캐시 조정¶

각 (pp_stage, tp_rank) 슬롯은 자체 레이어별·랭크별 KV 샤드를 소유합니다. 시퀀스의 입장(admission)은 그리드 전체에서 정합적이어야 합니다. 어느 한 랭크라도 시퀀스를 수용할 수 없다면, 스테이지는 일부 랭크만 결과를 캐싱한 채 샤딩된 forward 패스를 수행하게 되어 조용히 잘못된 결과를 생성합니다.

코디네이터는 coordinated_2d_admission 을 사용합니다. 이 함수는 두 단계 체크를 수행합니다. 모든 슬롯에 대한 드라이런 이후, 모든 슬롯이 수락한 경우에만 커밋합니다. 거부 시 어느 슬롯도 수정되지 않습니다.

use mlxcel::distributed::pipeline::{
    coordinated_2d_admission,
    CacheAdmissionRequest,
    PpTpAdmissionOutcome,
    PpTpCoord,
};

let mut grid = vec![
    (PpTpCoord::new(0, 0), &mut mgr_s0_r0),
    (PpTpCoord::new(0, 1), &mut mgr_s0_r1),
    (PpTpCoord::new(1, 0), &mut mgr_s1_r0),
    (PpTpCoord::new(1, 1), &mut mgr_s1_r1),
];
let req = CacheAdmissionRequest::new(42, prompt_len);
match coordinated_2d_admission(&mut grid, &req)? {
    PpTpAdmissionOutcome::Admitted => {
        // 스테이지 정합성이 보장되었으므로 계속 진행합니다.
    }
    PpTpAdmissionOutcome::Rejected { at, reason } => {
        tracing::warn!(%at, ?reason, "2D admission rejected");
    }
}

클러스터 TOML (2×2 예시)¶

[cluster]
name = "llama-70b-2x2"
pipeline_parallel_size = 2
tensor_parallel_size = 2
transport_backend = "rdma"

# 스테이지 0, TP 랭크 0
[[nodes]]
id = "s0-r0"
address = "192.168.1.10:8080"
role = "pipeline_tensor_parallel"
stage = 0
rank = 0

# 스테이지 0, TP 랭크 1
[[nodes]]
id = "s0-r1"
address = "192.168.1.11:8080"
role = "pipeline_tensor_parallel"
stage = 0
rank = 1

# 스테이지 1, TP 랭크 0
[[nodes]]
id = "s1-r0"
address = "192.168.1.20:8080"
role = "pipeline_tensor_parallel"
stage = 1
rank = 0

# 스테이지 1, TP 랭크 1
[[nodes]]
id = "s1-r1"
address = "192.168.1.21:8080"
role = "pipeline_tensor_parallel"
stage = 1
rank = 1

CLI 진입점¶

mlxcel generate \
  --model ./models/llama-70b-4bit \
  --pp-size 2 \
  --tp-size 2 \
  -p "Hello" -n 50

src/commands/generate.rs의 검증기는 이 조합을 받아들입니다. 일관성 없는 토폴로지(예: 방출된 클러스터 TOML의 노드 수가 pp_size × tp_size보다 작은 경우)에 대해 가독성 있는 에러를 방출합니다.

패리티 보고서¶

작은 Llama 계열 프록시로 --pp-size 2 --tp-size 2 실행이 고정 시드에서 단일 디바이스 레퍼런스와 동일한 그리디 디코딩 토큰을 생성함을 검증합니다. 이 검사를 2D PP × TP 토큰 동치성의 기준 검증으로 사용하십시오.

PP와 LoRA¶

파이프라인 병렬화는 LoRA 어댑터와 합성이 가능합니다. 단일 프로세스 어댑터 추론을 구동하는 --adapter 플래그가 파이프라인 스테이지 전반의 어댑터 추론도 동일하게 구동하며 같은 어댑터 디렉터리 레이아웃과 같은 랭크 / 스케일 의미를 그대로 사용합니다. 여러 스테이지에 걸쳐 어댑터를 실행하려고 별도 툴체인이나 별도 어댑터 포맷을 준비할 필요가 없습니다. 현재 구현은 Llama 패밀리에서 단일 어댑터 합성을 지원하며 다중 어댑터 스태킹과 런타임 핫스왑은 지원하지 않습니다.

스테이지 로컬 어댑터 로딩¶

각 파이프라인 스테이지는 자신에게 할당된 레이어 범위에 대한 기본 모델 가중치를 로드한 뒤, 스테이지 실행기가 실체화되기 전에 해당 기본 가중치에 LoRA 델타를 합성합니다. 스테이지별 로더는 mlxcel-core의 필터링된 로더로 어댑터의 adapters.safetensors(혹은 fallback filename인 adapter_model.safetensors) 파일을 스캔하며 레이어 인덱스가 스테이지 할당 범위를 벗어나는 어댑터 텐서를 스킵합니다. 범위를 벗어난 어댑터 텐서는 스테이지 Rust 측 가중치 맵에 아예 들어오지 않으므로 버퍼링도 폐기도 일어나지 않고 힙 트래픽과 MLX 측 배열 구성을 모두 절약할 수 있습니다.

퓨전 자체는 비-PP 경로와 공유됩니다. 두 경로 모두 동일한 fuse_lora_weights_into 프리미티브를 호출하므로, LoRA 식 W_fused = W_base + scale * (lora_b @ lora_a) 연산의 구현은 한 곳에만 존재하고 mlx-lm 대 HuggingFace PEFT 모양 규약을 협상하는 지점도 한 곳에만 존재합니다.

graph LR
    A[adapter_config.json] --> C[AdapterConfig]
    B[adapters.safetensors] --> D[load_stage_adapter_weights]
    C --> E[apply_stage_lora_adapter]
    D --> E
    S[스테이지 레이어 범위] --> D
    F[스테이지의 기본 가중치 맵] --> E
    E --> G[퓨전된 스테이지 가중치 맵]
    G --> H[스테이지 실행기]

CLI 사용법¶

# Llama 체크포인트에서 2-스테이지 PP + LoRA.
mlxcel generate \
  --model ./models/llama-3.2-1b-4bit \
  --adapter ./models/llama-3.2-1b-4bit-lora \
  --pp-size 2 \
  -p "Hello" \
  -n 64 \
  --temp 0

같은 플래그와 어댑터 디렉터리 레이아웃은 인-프로세스 파이프라인 병렬 런타임이 설정된 mlxcel-server에서도 동일하게 동작합니다.

동치성 계약¶

고정된 프롬프트, 고정된 시드, --temp 0 조건에서 PP + LoRA 실행의 디코드 토큰 스트림은 단일 프로세스 --adapter 실행과 비트 단위로 일치해야 합니다. 이 계약은 무시(ignored) 통합 테스트 tests/pipeline_lora_composition_real_models.rs 에서 검증합니다. 해당 테스트는 Llama 체크포인트 위에 배포된 LoRA를 두 스테이지에 걸쳐 실행하고 단일 프로세스 어댑터 출력과의 엄격한 동일성을 단언합니다.

src/distributed/pipeline/partial_loading.rs 에 있는 단위 테스트는 다음을 독립적으로 검증합니다.

• 스테이지 로컬 어댑터 텐서 필터(레이어 범위 포함, 임베딩과 lm_head 소유권, language_model.model.layers. VLM 스타일 프리픽스 포함);
• load_stage_adapter_weights가 다른 스테이지 소유 어댑터 텐서를 버퍼링-후-폐기가 아니라 스킵한다는 점;
• 스테이지 필터된 기본 가중치 + 스테이지 필터된 어댑터 합성 경로가 해당 스테이지 레이어 범위에서 전체 퓨전 후 필터링 경로와 동일한 퓨전 결과를 생성한다는 점.

적용 범위와 가드¶

어댑터 지원을 가장 먼저 추가한 스테이지 실행기는 LlamaStageExecutor입니다. 다른 패밀리는 가드 (ensure_no_adapter)로 라우팅되며 PP 스테이지 실행기에서 합성 단계를 구현하지 않은 모델 패밀리에 어댑터가 전달되면 소리 없이 어댑터가 무시되는 대신 명시적으로 실패합니다. 새 패밀리를 지원 매트릭스에 추가하려면 3줄 정도의 변경이면 됩니다. 패밀리의 load 시그니처에 adapter_path: Option<&Path>를 받고 가중치 로드와 filter_weight_map 호출 사이에서 apply_stage_lora_adapter를 호출한 뒤, 가드에서 해당 패밀리를 제거하면 됩니다.

현재 범위에서 제외되는 항목:

• PP 클러스터 위에서 여러 어댑터를 겹쳐 적용하는 다중 어댑터 스태킹. 스테이지 초기화 시점에 단 하나의 어댑터만 로드 가능합니다.
• 활성 PP 클러스터 위에서 어댑터의 런타임 핫스왑. 어댑터는 스테이지 초기화 중 한 번만 해석 및 퓨전됩니다.
• 텐서 병렬화 + LoRA + PP. 기존 TP 검증기는 여전히 어댑터 경로를 거부합니다. 2D TP x PP 합성에서 어댑터 지원은 현재 구현에 포함되지 않습니다.
• 원격 코디네이터 PP + 어댑터. 원격 파이프라인 토폴로지는 현재도 어댑터 경로에서 명시적으로 실패합니다. 어댑터 합성은 인-프로세스 서버 런타임과 CLI mlxcel generate 경로를 대상으로 합니다.

파이프라인 검증 범위¶

멀티 호스트 파이프라인 병렬 동작은 두 단계의 검증으로 보호됩니다. 이 섹션은 운영자가 기대할 수 있는 보장만 요약합니다. 정확한 워크플로 파일, 스크립트 이름, 테스트 모듈 경로는 유지보수자용 정보이므로 사용자 매뉴얼에서는 다루지 않습니다.

논리적 검증 단계는 모델 가중치 없이 실행되도록 설계되어 있습니다. 자동 파티셔너가 메모리 제약이 있는 스테이지에 더 적은 레이어를 배정하는지, 그리고 어느 스테이지라도 KV 캐시 용량이 부족하면 긴 프롬프트 입장을 일관되게 거부하는지 검증합니다. 이는 나중에 OOM이나 스테이지 간 캐시 불일치로 드러날 수 있는 과도한 입장을 방지합니다.

실모델 검증 단계는 전용 멀티 호스트 하드웨어와 로컬 모델 가중치를 사용하므로 명시적으로 요청한 경우에만 실행됩니다. 전송 동작, 활성화 핸드오프, 스테이지 스케줄링, 원격 파이프라인 클러스터의 토큰 동치성에 영향을 줄 수 있는 변경 전에 사용하는 검증 단계입니다.

탄력적 리파티셔닝 (실험적)¶

기본적으로 파티션 플랜은 클러스터 시작 시점에 고정됩니다. opt-in 탄력적 리파티셔닝 플로우를 사용하면 클러스터를 전체 재시작하지 않고도 실행 중에 스테이지 간 레이어 분포를 다시 계산할 수 있습니다.

플로우 활성화¶

코디네이터 프로세스의 mlxcel-server에 --enable-elastic-pp를 전달합니다. 이 플래그는 순수하게 additive이며 플래그를 지정하지 않은 클러스터는 기존 동작을 그대로 유지합니다.

mlxcel-server -m models/llama-3.2-1b-4bit \
  --distributed-config cluster.toml --node-id coordinator \
  --parallel 2 --pp-micro-batch-size 2 \
  --enable-elastic-pp \
  --elastic-pp-drain-timeout 120 \
  --elastic-pp-pressure-fraction 0.92 \
  --elastic-pp-cool-down 30

트리거¶

현재 두 가지 트리거를 지원합니다:

부하 예측에 기반한 예측적 정책은 지원하지 않습니다.

드레인 시맨틱¶

트리거가 수락되면 코디네이터는 다음 상태 머신을 따라 각 단계를 진행합니다:

Idle → Draining → Rebalancing → Resuming → Idle

1. Draining: 코디네이터가 모든 스테이지에 BeginDrain을 팬 아웃하고 드레인 플래그가 설정된 동안 서버는 새 시퀀스 입장을 패닉(503 대응)으로 거부합니다. 각 스테이지의 in_flight_requests == 0 보고를 probe_stages로 폴링하며 타임아웃(--elastic-pp-drain-timeout)을 초과하면 실패합니다.
2. Rebalancing: 바이트 정확 밸런서 (partition_balance::balance_layers)를 최신 스테이지별 메모리 버짓과 함께 호출합니다. 기존 인접성 제약(Gemma 4, Jamba)을 그대로 재사용하며 새 플랜은 partition::validate_partition으로 검증합니다.
3. Resuming: 각 스테이지가 새 LayerFilter로 partial_loading 경로를 재진입하여 레이어 가중치를 다시 로드합니다. KV 캐시는 처음부터 다시 빌드됩니다. 드레인된 시퀀스는 클라이언트가 HTTP 재시도로 다시 제출합니다. 드레인이 리밸런싱 이전에 완료되었기 때문에 사용자 작업은 유실되지 않습니다.

어느 단계에서든 실패가 발생하면 코디네이터는 터미널 상태 Failed로 전환합니다. 이전 레이어 범위가 계속 활성 상태를 유지하고 드레인은 best-effort로 해제되어 서비스를 지속할 수 있습니다. 재시도하려면 다른 explicit 트리거가 필요합니다.

관측 가능성¶

상태가 전이될 때마다 구조화된 RepartitionEvent가 방출되어 메트릭 서브시스템으로 전달됩니다. 구체적인 메트릭 이름과 /metrics, --debug-pp-trace에서 노출되는 트레이스 이벤트는 아래 Observability 섹션을 참조하십시오.

제약¶

• 진행 중 요청을 유지한 채 크로스 노드 레이어 마이그레이션은 지원하지 않습니다. 현재 플로우는 선 드레인 후 리파티셔닝 방식입니다. 진정한 zero-downtime 마이그레이션은 현재 구현에서 지원하지 않습니다.
• 피어 간 가중치 스트리밍은 범위에서 제외합니다. 스테이지는 디스크의 원본 safetensors에서 재로드합니다.
• 파이프라인 축만 탄력적으로 재구성합니다. 2D (pp_stage, tp_rank) 양축의 동시 리파티셔닝은 지원하지 않습니다.
• HTTP 관리자 엔드포인트는 없습니다. explicit 트리거는 코디네이터 인프로세스 API로 전달합니다.

관측 가능성 (Observability)¶

Prometheus /metrics 엔드포인트는 파이프라인 병렬화를 위한 네 가지 메트릭 패밀리와 오프라인 분석용 chrome-tracing JSON 라이터, 그리고 귀책 정보가 담긴 OOM 진단 기능을 제공합니다.

/metrics 엔드포인트¶

mlxcel-server --metrics 또는 mlxcel-server --metrics-port <PORT>로 활성화합니다. 엔드포인트는 현재 메인 HTTP 리스너에 멀티플렉싱됩니다. 별도 포트를 지정하면 경고를 출력합니다. 엔드포인트는 기존의 서버 레벨 요청 카운터에 더해 다음 파이프라인 병렬 패밀리를 노출합니다:

reason은 memory, sequence_cap, already_cached, rejected_by_stage 중 하나입니다. trigger는 explicit, memory_pressure 중 하나이며 outcome은 completed, aborted, failed, 또는 progress(중간 상태 전이)입니다.

Chrome-tracing 모드 (--debug-pp-trace)¶

mlxcel-server --debug-pp-trace <path>는 PpTracer를 활성화합니다. 모든 스케줄러 액션이 메모리에 버퍼링되며 최종 JSON 파일은 정상 종료(또는 명시적 플러시) 시점에 기록됩니다. chrome://tracing 또는 https://ui.perfetto.dev에서 파일을 열어 타임라인을 확인합니다.

기록되는 이벤트:

• batch_arrival: 스케줄러가 새 요청 배치를 수신합니다.
• stage (duration): 스테이지 포워드 패스를 감쌉니다. tid 컬럼은 stage_index로 그룹핑됩니다.
• activation_send / activation_recv: 각 핸드오프마다 한 쌍이 기록되며 bytes와 방향 인자를 포함합니다.
• admission_reject: reason, required_bytes, available_bytes를 포함합니다.

트레이스 파일 기록은 idempotent합니다. 최초로 성공한 플러시의 결과만 유지되며 Drop 시에는 다시 기록하지 않습니다.

OOM 진단¶

KV 캐시 입장 거부는 coordinated_admission_with_attribution으로 거부한 스테이지 식별자와 KV 점유를 노출합니다. 반환되는 AdmissionDiagnostic은 다음을 포함합니다:

• stage_index: 어떤 스테이지가 거부했는지.
• reason: enum variant (Prometheus에서 사용하는 저카디널리티 단축 이름은 metric_label()로 얻습니다).
• used_memory_bytes, budget_bytes: 거부 시점의 해당 스테이지 KV 회계.
• active_sequences: 현재 점유 시퀀스 수.

서버 에러 경로는 일반적인 500 대신 이 진단을 사람이 읽을 수 있는 메시지로 포맷팅합니다.

Grafana 대시보드¶

이 메트릭을 위한 일반적인 Grafana 대시보드는 다음 항목을 포함합니다:

1. 스테이지별 busy fraction (timeseries).
2. 평균 버블 비율 (빨간색/주황/녹색 임계치가 있는 stat 패널).
3. 스테이지 쌍별 활성화 전송 지연 p95, p99.
4. (stage, reason)별 거부율.
5. (trigger, outcome)별 누적 리파티션 이벤트.

mlxcel-server --metrics를 스크레이핑하는 Prometheus 데이터 소스를 대상으로 임포트합니다. Grafana 스택 자체는 오퍼레이터가 준비하며 본 배포에서는 대시보드 JSON만 제공합니다.