캡처로 JAX OOM 디버깅하기
개요
jit 컴파일된 JAX 단계는 모든 스크래치 메모리를 하나의 계획된 임시 블록으로 할당받습니다. 따라서 단계의 메모리는 XLA의 계획이며, .xplane.pb 파일에는 해당 계획과 할당자의 런타임 이벤트가 모두 포함됩니다. scripts/oom_report.py 스크립트는 표준 라이브러리만 사용하여 두 정보를 2-5초 내에 읽어들입니다.
- 계획: 각 모듈의 버퍼 할당 정보로,
/host:metadata에서 가져옵니다. 파라미터(얼마나 기부되었는지 포함), 별칭이 없는 출력, 임시 블록, 힙 피크(heap peak)를 제공합니다. 피크 시점에 활성화된 버퍼의 크기, 형상, 역할, Python 코드 라인을 보여줍니다. 역할은 JAX 이름 스택에서 비롯됩니다:jvp(는 포워드(역전파를 위해 유지),transpose(는 백워드,s8사용자 정의 호출 출력은 라이브러리 워크스페이스입니다. - 런타임: BFC 할당자의 이벤트입니다. 실행 시작 시의 용량과 사용 중인 바이트, 피크,
AllocatorRetry대기, 커널을 실행하지 않은 실행, 마지막 두 실행 사이의 호스트 이벤트를 제공합니다. 또한 마지막 실행 시작 시 활성화된 버퍼를 할당한 호스트 이벤트별로 그룹화하여 보여줍니다:jit_step execution 3의 출력,BufferFromHostBuffer(s32[…])(device_put), 또는 "캡처 시작 전에 할당됨".
보고서는 각각 측정 또는 추정으로 태그된 진단 줄로 끝납니다. 먼저 실행하세요:
python3 -P <이 스킬의 기본 디렉토리>/scripts/oom_report.py <캡처 디렉토리 또는 .xplane.pb> [--module 'jit_train_step(123)'] [--top 25] [--json out.json]
python3 -P <이 스킬의 기본 디렉토리>/scripts/oom_report.py --self-test
기본적으로 마지막 실행이 커널을 시작하지 않은 모듈을 분석합니다. 해당하는 모듈이 없으면 실행된 가장 큰 모듈을 사용합니다. --module을 사용하여 다른 모듈(예: 평가 단계)을 선택할 수 있습니다.
진단 줄에서 원인 및 수정 방안으로
| 진단 줄 | 원인 | 제안할 수정 방안; 절감량은 보고서의 바이트 수 기준 |
|---|---|---|
단계 자체가 맞지 않음 | 계획만으로 풀이 초과됨 | 피크 시점의 역할 분할 참고(아래) |
↳ 주로 포워드, 역전파를 위해 유지됨 | 활성화: 배치 × 시퀀스에 따라 증가; naive 어텐션은 레이어당 B·H·S·S 추가 | jax.checkpoint / 블록에 remat(포워드 유지 바이트 대부분 절감, 약 한 번의 포워드 패스 비용 증가); 그래디언트 누적과 함께 더 작은 마이크로 배치; 더 짧은 시퀀스(S² 버퍼가 S 절반일 때 4배 감소). 융합 어텐션은 "B·H·S·S 버퍼가 계획을 떠나는지 확인"으로만 제안 |
↳ 주로 백워드 | 그래디언트 및 임시 변수 | 샤딩(FSDP) 또는 누적; bf16 그래디언트의 f32 복사본 확인 |
↳ 워크스페이스 | GEMM / conv 라이브러리 스크래치 | 일반적으로 작음; ROCm에서는 gfx950에서 hipBLASLt GEMM당 64MiB |
실행당 X MiB 증가, = 비별칭 출력의 1.00배 | 프로그램이 각 단계의 출력을 활성 상태로 유지(단계 간 누수): 활성 목록에 ... execution 1, 2, ...의 출력 표시 | 실행 간 호스트 이벤트에서 보유자를 찾음(예: $<unknown> append = list.append). 그런 다음 참조 해제, 단계 내에서 축소(트리가 아닌 노름), 또는 jax.device_get 실행. 더 큰 풀은 OOM을 지연시킬 뿐 |
기부된 파라미터 없음 ... 최대 X까지 절감 가능 | 이전 및 새 파라미터/옵티마이저 상태가 함께 활성 상태 | 상태 인수에 donate_argnums 사용 |
다른 활성 배열이 공간을 차지함; 가장 큰 것은 ...에서 비롯됨 | 단계 외부의 배열: GPU에 배치된 데이터세트 또는 평가 복사본, 프리페치, 기타 jit 출력 | 호스트에 유지하고 스트리밍, 해제, 또는 dtype 축소. "캡처 시작 전에 할당됨"은 이름을 지정하기 위해 더 일찍 프로파일링해야 함 |
AllocatorRetry ... 가장 긴 N 초, 하나도 시작되지 않음 | OOM이 캡처 내에서 발생 | 증거만 해당. 실패한 요청은 일반적으로 해당 모듈의 임시 블록 또는 하나의 출력 |
XLA가 이미 N개의 명령어를 재구성화함 | XLA 자체의 remat이 이미 한계에 도달함 | 수정은 XLA가 아닌 프로그램에서 비롯될 것으로 예상 |
각 수정 방안에 대해 추정된 절감량과 비용(추정으로 레이블 지정)을 제시하세요. 추정치를 용량과 비교하여 확인하세요. "XLA_PYTHON_CLIENT_MEM_FRACTION 높이기"는 GPU에 여유 공간이 있는 경우에만 맞지 않는 계획을 수정합니다. 용량이 HBM보다 훨씬 낮으면 비율이 설정되었음을 의미합니다. 캡처는 환경 변수를 기록하지 않습니다.
JAX 용어로 표현된 플레이북의 메모리 해부학
Ultra-Scale Playbook("메모리 사용량 프로파일링")은 훈련 메모리를 가중치, 그래디언트, 옵티마이저 상태 및 활성화로 나눕니다. 보고서에서:
- 가중치 + 옵티마이저 상태 = 단계의 파라미터. fp32의 Adam은 파라미터당 12B(p, m, v). 기부됨(donated)은 복사본이 없음을 의미.
- 그래디언트 = 피크 시점의
backward역할, 단계가 반환하는 경우 출력도 포함. - 활성화 =
forward, kept for the backward pass. Scan /nn.scan모델은while루프에[layers, …]형식의 누적 버퍼로 유지. - "1단계는 맞지만 N단계는 실패": PyTorch에서는 첫 단계에서 생성된 옵티마이저 상태가 원인. JAX에서는 상태가 1단계 이전에 존재하므로 대신 증가 라인을 찾으세요. 첫 번째 호출은 컴파일 및 자동 튜닝도 수행하므로 첫 번째 실행 전의 피크는 컴파일 시간일 수 있습니다.
규칙
- 캡처만 사용하세요. 수정 방안을 테스트하기 위해 모델을 다시 빌드하거나 워크로드를 다시 실행하지 마세요. 계획의 바이트 수에서 추정하고 사용자 측에서 확인 방법(
jit(...).lower(...).compile().memory_analysis())을 알려주세요. - 캡처 디렉토리에 쓰지 마세요. XProf의 변환기(
xspace_to_tool_names,memory_viewer,xprofCLI)는.xplane.pb옆에*.hlo_proto.pb캐시 파일을 씁니다. XProf 뷰가 필요하면 먼저.xplane.pb를 임시 디렉토리에 복사하세요. - 모든 숫자에 레이블을 지정하세요: 측정됨, 읽음 또는 추정됨.
측정된 참조 (단일 MI355X, 소규모 풀, 2026-11-06)
| 캡처 | 보고서에 표시되는 내용 |
|---|---|
| 어텐션 계획 너무 큼 | 3.72 GiB 필요 > 2.79 GiB; 힙 피크 81.6% 포워드 유지; f32[8,8,1024,1024] ×5 소프트맥스 라인; 임시 블록 3,495 MiB = 오류의 "3.41GiB"; AllocatorRetry 10초; 89개의 remat 명령어 |
| 모든 단계에서 그래디언트 유지 | 단계 시작 시 사용 중 316 → 1,252 MiB, 단계당 +104.0 MiB = 비별칭 출력의 1.00배; 계획 788 MiB 적합; 단계 간 $<unknown> append |
| 기부 없음 | 3.88 GiB 필요 > 2.79 GiB; 1,902 MiB 파라미터 + 1,902 MiB 출력, 기부 0; 임시 블록 172 MiB에 불과 |
| MaxText llama2-7b, 8 GPU, FSDP (적합) | 임시 61.9 GiB; 힙 피크 57% 포워드 유지, 스캔 스택 bf16[32,4,4096,11008] ×2 via while.50 |
제한 사항
- 파일당 하나의 호스트: 멀티 호스트 실행은 호스트당 하나의
.xplane.pb파일을 가지므로 각각 실행해야 합니다. - HBM만 해당: 메모리 공간 0만 분석됩니다. 호스트 오프로드는
memory space 5로 표시됩니다. - 연속성: "적합"은 바이트를 비교합니다. BFC는 또한 연속적인 청크가 필요하므로 근접 적합도 단편화로 인해 실패할 수 있습니다.
- "other" 역할: JAX 이름 스택이 없는 XLA 생성 연산으로 정의된 버퍼와 해당 스택이 있는 별칭이 없는 버퍼.
- CUDA: 시도되지 않았습니다. 동일한 BFC 및 XLA 이벤트를 사용하지만 이름은 확인되지 않았습니다.
- TPU / eager: TPU 캡처는 힙 피크 및 실행 없이 계획 총합만 제공합니다. Eager(연산별) 프로그램은 많은 소규모 모듈로 표시되므로 런타임 섹션을 참조하세요.