[gspark-etri]

배경

v4.1.1(62a7831)에서 아래 세 결함이 재현됩니다. 세 결함 모두 optimus-rm에서는 이미 고쳤고(#187, #188, #189), 이 PR은 그 세 커밋을 main에 그대로 cherry-pick한 것입니다(cherry picked from commit ... 표시 포함). main 사용자를 위한 v4.1.2 핫픽스용입니다.

재현과 검증은 v4.1.1이 명시한 지원 조합(opt_prime/README.md: PyTorch 2.5.0 / CUDA 12.4, transformers==4.46.2)에서 했습니다. 환경은 Python 3.10, torch 2.5.0+cu124, transformers 4.46.2, A100 80GB 한 장입니다.

1. KV 캐시 생성이 배치 크기 2 이상에서 크래시 (#164, 원본 d46b8ff)

문제

• Optimus_Inference(..., use_kv_cache=True, dynamo_capture=True)로 generate()를 부르면 배치 크기 1에서는 HF greedy와 같은 결과가 나옵니다.
• 배치 크기 2에서는 prefill 단계에서 RuntimeError: mat1 and mat2 shapes cannot be multiplied (2x128 and 64x64)로 크래시합니다.
• 즉 v4.1.1에서 KV 캐시 생성은 사실상 배치 크기 1에서만 동작합니다.

원인

• prefill에서 자동으로 만드는 position_ids의 모양이 [1, S]입니다. 그런데 torch.export로 캡처한 그래프는 position_ids의 배치 차원을 input_ids([B, S])와 같은 크기로 묶어 둡니다.
• decode에서는 position_ids를 [[position]]([1, 1])로 만듭니다.
• 다음 토큰 입력은 next_token.unsqueeze(0)이라 [B]가 [B, 1]이 아니라 [1, B]가 됩니다.
• B=1일 때는 이 차이가 드러나지 않습니다.

수정 (inference.py, schedule_infer.py, +5/−7줄)

• prefill: torch.arange(seq_len).expand(B, -1)로 [B, S]를 만듭니다.
• decode: torch.full((B, 1), position)으로 [B, 1]을 만듭니다.
• 다음 토큰: unsqueeze 분기 세 줄을 next_token.reshape(-1, 1) 한 줄로 바꿉니다.

테스트: tests/test_kv_cache_batch.py

• 작은 Llama(2층, hidden 64)로 KV 캐시를 켜고, B=1과 B=2 각각 HF generate() greedy 출력과 토큰 단위로 같은지 비교합니다.
• GPU가 없으면 skip합니다.

2. TP/DP rank의 체크포인트가 서로 덮어씀 (#165, 원본 fb09a61)

문제

• save_ckpt()는 파일 이름을 pipeline stage로만 정합니다(ckpt_step_{step}_epoch_{epoch}.pt).
  • tp=2이면 두 TP rank가 서로 다른 shard를 같은 파일에 써서, 한 shard가 사라집니다.
  • 다시 불러오면 두 rank가 같은 shard를 로드해 잘못된 가중치로 학습을 이어갑니다.
• dp>1이면 모든 DP replica가 같은 파일을 동시에 쓰고, 오래된 파일도 동시에 지웁니다. 이때 os.remove가 FileNotFoundError를 낼 수 있습니다.
• save_lora_ckpt()도 DP replica마다 저장합니다.
• ckpt_dir_postfix를 주지 않으면 디렉터리 이름이 checkpoint_None_<stage>가 됩니다.
  • 원인: f"checkpoint_{postfix}_{stage}" or f"checkpoint_{stage}"에서 왼쪽 문자열이 항상 참이라, or 오른쪽 기본값이 한 번도 쓰이지 않습니다.

수정 (opti_pri.py, +13/−4줄)

• save_ckpt와 save_lora_ckpt는 DP rank 0만 저장합니다((rank // tp) % dp == 0). 이미 save_hf_ckpt가 쓰는 방식입니다.
• tp>1이면 파일 이름에 _tp{rank % tp}를 붙입니다. 오래된 파일을 지울 때도 자기 TP shard의 파일만 지웁니다.
• load_ckpt는 같은 이름 규칙으로 읽습니다.
• postfix가 없으면 디렉터리 이름은 checkpoint_<stage>입니다.

테스트: tests/test_ckpt_names.py

• tp=2, dp=2 rank를 흉내 내어 다음을 확인합니다.
  • TP shard마다 파일이 하나씩 생기는지
  • DP rank 0만 쓰는지
  • LoRA를 다시 불러올 때 올바른 shard를 읽는지
  • PP만 쓰는 경우 기존 파일 이름이 그대로인지
  • 디렉터리 이름이 맞는지

3. MPS 매니저가 같은 사용자의 다른 GPU 작업 데몬을 죽임 (#166, 원본 0e69cc9)

문제

• _mps_daemon_running()은 우리 pipe 디렉터리에 control 소켓이 없어도, 호스트 전체에서 pgrep nvidia-cuda-mps-server/control이 하나라도 걸리면 "실행 중"으로 판단합니다.
• 그러면 다른 작업이 자기 pipe 디렉터리로 띄운 데몬을 우리 것으로 오인합니다. 이후 우리 pipe 디렉터리로 연결을 시도하다 실패하면 "stale 데몬"으로 보고 _force_cleanup_stale_daemon()을 부르고, 이 함수가 pkill -u $UID로 그 사용자의 MPS server와 control 프로세스를 전부 죽입니다.
• 예: 작업 A가 /tmp/mps-a로 MPS를 쓰는 중에 작업 B가 기본 디렉터리로 시작하면, 작업 A의 MPS 데몬이 죽습니다.

수정 (mps_manager.py, +8/−38줄)

• 데몬 실행 여부는 우리 pipe 디렉터리의 control 소켓만으로 판단합니다. 호스트 전체를 보는 pgrep 검사를 지웁니다.
• stale 정리 단계에서 pkill 루프를 지웁니다. 우리 pipe 디렉터리에 quit를 보내고, 우리 소켓만 지웁니다.

테스트: tests/test_mps_manager.py

• 가짜 pipe 디렉터리를 두고, subprocess.run을 바꿔치기해 "다른 작업의 데몬이 떠 있는" 상황을 만듭니다.
• 그 상황에서 우리 데몬이 실행 중으로 판단되지 않는지, pkill이 호출되지 않는지 확인합니다.

검증

명령: cd opt_prime && TORCH_DEVICE_BACKEND_AUTOLOAD=0 CUDA_VISIBLE_DEVICES=<GPU> PYTHONPATH=$PWD python -m pytest -q -rA tests

환경: torch 2.5.0+cu124, transformers 4.46.2, A100 80GB 한 장.

	결과
수정 전: 62a7831에 테스트 파일만 추가	8 failed, 3 passed
수정 후: 이 브랜치	9 passed

• 수정 전 실패 내역
  • KV: torch.export 경로에서 B=2가 실패하고 B=1은 통과합니다. 아래 HFTracer 실패 2건이 더해집니다.
  • 체크포인트: 3건 실패
  • MPS: 2건 실패
• CPU 전용 환경(GPU 없음, torch 2.10.0+cpu): 7 passed, 2 skipped. KV 테스트만 skip됩니다.
• 실제 멀티 GPU TP/DP 학습과 실제 MPS 데몬으로는 시험하지 않았습니다. 단위 테스트와 GPU 한 장 검증까지만 했습니다.

호환성과 제한

• 체크포인트 디렉터리 이름 변경: postfix 없이 저장한 기존 체크포인트는 디렉터리가 checkpoint_None_<stage>입니다. 이 PR 이후에는 checkpoint_<stage>를 찾으므로, 그런 체크포인트로 재개하려면 디렉터리 이름을 바꿔야 합니다. PP만 쓰는 경우 파일 이름은 바뀌지 않습니다.
• TP 체크포인트 파일 이름 변경: tp>1이면 이제 _tp{rank}가 붙습니다. 그 전 파일은 shard 하나가 이미 사라진 상태라 원래부터 제대로 복원할 수 없었습니다.
• dp>1 재개 조건: 모든 replica가 같은 체크포인트 디렉터리를 볼 수 있어야 합니다.
• MPS 강제 종료 제거: 우리 데몬이 quit에 응답하지 않고 멈춰 있어도 더는 강제로 죽이지 않고, 새 데몬을 띄웁니다.
• 다루지 않은 것: v4.1.1에서 기본 경로인 HFTracer(dynamo_capture=False)로 KV 캐시를 켜면, 배치 크기와 상관없이 tracing 단계에서 NotImplementedError: argument of type: <class 'transformers.utils.fx.ProxyableDynamicCache'>로 실패합니다. 이 PR 전후 동작이 같은 별개 문제라 이 PR에서는 다루지 않습니다.
• main에 생기는 테스트: main에는 테스트 디렉터리와 CI가 없어서, opt_prime/tests에 이 세 테스트만 생깁니다.

관련 Issue

Refs #164, #165, #166 (optimus-rm에서 #187, #188, #189로 수정, 이 PR은 main 백포트)

🤖 Generated with Claude Code