[pytorch] CI: fbgemm/torchrec 핀 버전 업데이트 및 빌드 로직 리팩토링

PR 링크: pytorch/pytorch#172179 상태: Merged | 변경: +52 / -50

들어가며

PyTorch의 CI는 fbgemm(Facebook GEMM)과 torchrec(Torch Recommendations)을 특정 커밋에 고정(pin)하여 빌드합니다. 이 PR은 두 라이브러리의 핀 버전을 릴리즈 태그로 업데이트하고, ROCm과 CUDA 양쪽에서 중복되던 fbgemm 빌드 로직을 install_fbgemm 함수로 분리 리팩토링한 작업입니다.

핵심 코드 분석

1. 핀 버전을 릴리즈 태그로 전환

Before:

# fbgemm.txt
de731af65b4f04696e85c729e3282450b51b95fd

# torchrec.txt
6cd9fd362514d14ebb9ed51314c62ac1e1e2bbf2

After:

# fbgemm.txt
v1.5.0-release

# torchrec.txt
release/v1.4.0

커밋 해시 대신 릴리즈 태그를 사용하여 어떤 버전인지 즉시 파악할 수 있게 되었습니다.

2. install_fbgemm 함수 분리

기존에는 install_torchrec_and_fbgemm 함수 내에 ROCm용 fbgemm 빌드 로직이 인라인되어 있었고, CUDA용은 별도 경로를 탔습니다. 이를 install_fbgemm 함수로 분리했습니다.

Before:

function install_torchrec_and_fbgemm() {
  # ... torchrec 설치 ...
  if [[ "$BUILD_ENVIRONMENT" == *rocm* ]]; then
    # ROCm: 30줄의 fbgemm 빌드 로직 (wheel 캐시 확인, 빌드, 설치)
    git clone --recursive https://github.com/pytorch/fbgemm
    pushd fbgemm/fbgemm_gpu
    python setup.py bdist_wheel --build-variant=rocm
    # ...
  else
    # CUDA: 한 줄 pip install
    pip_build_and_install "git+https://github.com/pytorch/FBGEMM.git@${fbgemm_commit}#subdirectory=fbgemm_gpu"
  fi
}

After:

function install_fbgemm() {
  local build_variant=$1
  local fbgemm_commit=$(get_pinned_commit fbgemm)
  # wheel 캐시 확인 -> 없으면 빌드 -> 설치 (공통 로직)
  if [ "${found_whl}" == "0" ]; then
    git clone --recursive https://github.com/pytorch/fbgemm
    python setup.py bdist_wheel --build-target=default --build-variant="${build_variant}"
  fi
  pip_install_whl "${file}"
}

function install_torchrec_and_fbgemm() {
  # ... torchrec 설치 ...
  if [[ "$BUILD_ENVIRONMENT" == *rocm* ]]; then
    install_fbgemm "rocm"
  else
    install_fbgemm "cuda"
  fi
}

build_variant 인자 하나로 CUDA/ROCm을 구분하고, wheel 캐시 확인 및 빌드 로직을 공통화했습니다.

3. libtbb 워크어라운드

# TODO (huydhn): Newer FBGEMM has a bug in detecting libtbb when building from source
LIBTBB_PATH="$(find "$(dirname "$(which python)")/../lib/" -name libtbb.so.12)"
export LD_PRELOAD="$LIBTBB_PATH":"$LD_PRELOAD"

새 버전의 fbgemm에서 libtbb 탐지 버그가 있어, LD_PRELOAD로 강제 로드하는 워크어라운드가 추가되었습니다.

왜 이게 좋은가

빌드 로직의 함수화는 단순한 코드 정리를 넘어, ROCm과 CUDA 빌드 경로의 일관성을 보장합니다. 기존에는 ROCm에만 wheel 캐싱이 적용되었지만, 리팩토링 후에는 CUDA에서도 동일한 캐싱 로직이 적용됩니다. 또한 핀 버전을 릴리즈 태그로 전환하면 CI 실패 시 어떤 버전에서 문제가 발생했는지 즉시 파악할 수 있습니다.

정리

• fbgemm v1.5.0-release, torchrec release/v1.4.0으로 핀 업데이트
• fbgemm 빌드 로직을 install_fbgemm(build_variant) 함수로 분리
• CUDA/ROCm 양쪽에서 wheel 캐싱 로직 공통 적용
• libtbb 탐지 버그에 대한 LD_PRELOAD 워크어라운드 추가

참고 자료

• PR #172179
• fbgemm
• torchrec

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.