OCaml용 PyTorch 설정하기

2024년 11월 17일 · 태그: ocaml, cuda, python · 원문 · 한국어 번역(원문에서 변경됨)

OCaml은 견고하고 타입 안전한 시스템을 구축하는 데 강점이 있지만, 현대 AI 개발을 위한 OCaml 설정은 의외로 다루는 곳이 드문 주제입니다. 이 가이드는 대규모 언어 모델의 지속 사전 학습(continued pre-training)이나 전체 튜닝(full-tuning)을 수행하기 위해, H100 GPU 8개를 탑재한 NVIDIA HGX AI 플랫폼에서 OCaml 개발 환경을 구성하는 방법에 초점을 맞춥니다. 일부 퀀트 펀드가 AI 인프라에 OCaml을 도입하려고 많은 투자를 했지만, 이처럼 GPU가 밀집된 구성에서 OCaml을 설정하는 방법에 관한 공개 문서는 놀라울 만큼 적습니다. 여기서는 Ubuntu에서 CUDA 및 PyTorch와 함께 OCaml을 설정하는 전체 과정을 차례로 살펴봅니다. 특히 요즘 많은 펀드가 도입하고 있는 이러한 멀티 GPU 시스템만의 고유한 요구 사항에 주의를 기울입니다.

시스템 의존성

시작하기 전에 Ubuntu를 실행하고 있는지 확인하십시오. 이 설정은 Ubuntu에서 테스트되었으며, 몇몇 설치 단계에는 root 권한이 필요합니다.

# Check your Ubuntu version
lsb_release -a

먼저 기본 시스템 의존성을 설치합시다. 이 패키지들은 이후 설정에 필요한 필수 개발 도구와 라이브러리를 제공합니다:

sudo apt-get update
sudo apt-get install -y \
    build-essential \
    pciutils \
    wget \
    clang \
    libaio-dev \
    python3-pip \
    python3-venv

NVIDIA 드라이버 설치

NVIDIA 드라이버 설치 방법은 데스크톱 시스템을 실행하는지 헤드리스 시스템을 실행하는지에 따라 다릅니다:

# For desktop systems with X.org or Wayland
sudo apt-get -y install nvidia-driver-550

# For headless systems (servers without display)
sudo apt-get -y install --no-install-recommends nvidia-headless-550

# Install CUDA toolkit
sudo apt-get -y install nvidia-cuda-toolkit

이 명령은 NVIDIA 드라이버 550 버전과 CUDA 툴킷을 설치하며, CUDA 툴킷은 GPU 가속에 필요한 구성 요소를 제공합니다.

Python 환경 설정

먼저 pyenv를 설치하고 Python 3.12.4를 구성합시다:

# Install pyenv dependencies
sudo apt-get install -y make build-essential libssl-dev zlib1g-dev \
libbz2-dev libreadline-dev libsqlite3-dev wget curl llvm \
libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev libffi-dev liblzma-dev

# Install pyenv
curl https://pyenv.run | bash

# Add pyenv to your shell configuration (~/.bashrc or ~/.zshrc)
echo 'export PYENV_ROOT="$HOME/.pyenv"' >> ~/.bashrc
echo 'command -v pyenv >/dev/null || export PATH="$PYENV_ROOT/bin:$PATH"' >> ~/.bashrc
echo 'eval "$(pyenv init -)"' >> ~/.bashrc

# Reload shell configuration
source ~/.bashrc

# Install Python 3.12.4
pyenv install 3.12.4

# Set Python 3.12.4 as global version
pyenv global 3.12.4

# Verify installation
python --version  # Should output: Python 3.12.4

# Upgrade pip
python -m pip install --upgrade pip

# Install PyTorch with CUDA support
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

# Install NVIDIA cuDNN
pip install nvidia-cudnn-cu11

# Install Hugging Face libraries
pip install transformers datasets accelerate

# Install scientific computing packages
pip install scipy numpy pandas

# Install optimized attention implementations
pip install flash-attn --no-build-isolation
pip install -v -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers

# Install GPU monitoring tools
pip install nvitop

DeepSpeed 및 Megatron-LM 설치

DeepSpeed 설치

DeepSpeed는 대규모 모델 학습을 위한 최적화 기능을 제공합니다. CUDA 연산 지원을 포함하여 설치합시다:

# Basic installation
pip install deepspeed

# OR: Install from source for latest features
git clone https://github.com/microsoft/DeepSpeed.git
cd DeepSpeed
pip install -e .

# OR: Install with pre-compiled operations (recommended)
DS_BUILD_OPS=1 pip install deepspeed --global-option="build_ext" --global-option="-j8"

설치를 검증하고 사용 가능한 연산을 확인하십시오:

ds_report

고급 설치 옵션:

# Install specific operations
DS_BUILD_FUSED_ADAM=1 pip install deepspeed    # FusedAdam support
DS_BUILD_FUSED_LAMB=1 pip install deepspeed    # FusedLamb support
DS_BUILD_SPARSE_ATTN=1 pip install deepspeed   # Sparse attention support

# Skip CUDA version check if needed (use with caution)
DS_SKIP_CUDA_CHECK=1 pip install deepspeed

Megatron-LM 설치

Megatron-LM을 사용하면 대규모 트랜스포머 모델을 대규모로 학습할 수 있습니다. 설정 방법은 다음과 같습니다:

# Install NVIDIA Apex first
git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --disable-pip-version-check --no-cache-dir \
    --global-option="--cpp_ext" \
    --global-option="--cuda_ext" ./
cd ..

# Install Megatron-LM
git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM
git checkout core_r0.5.0
pip install --no-use-pep517 -e .

OCaml 설정

마지막으로 PyTorch 바인딩과 함께 OCaml 환경을 설정합니다:

# Install OCaml package manager
sudo apt-get install -y opam m4 pkg-config

# Initialize OPAM
opam init --auto-setup --yes
eval $(opam env)

# Create new OCaml environment
opam switch create 4.14.1
eval $(opam env)

# Install OCaml development tools
opam install -y dune merlin ocaml-lsp-server odoc ocamlformat utop

# Install Torch dependencies and bindings
opam install -y ctypes ctypes-foreign conf-pkg-config
opam install -y torch

이 과정에서 설정되는 항목은 다음과 같습니다:

OPAM(OCaml 패키지 관리자)
OCaml 4.14.1 환경
다음을 포함한 개발 도구:
- Dune(빌드 시스템)
- Merlin(코드 자동 완성)
- LSP 서버 (IDE 통합)
- 문서화 도구
- UTop (향상된 REPL)
OCaml용 PyTorch 바인딩

검증

설치하고 시스템을 재부팅한 후 설정을 검증하십시오:

# Check NVIDIA driver installation
nvidia-smi

# Check GPU interconnect topology
nvidia-smi topo -m

# Start OCaml REPL and test Torch
utop
# In utop:
open Torch;;

nvidia-smi 출력에 GPU가 표시되어야 하며, OCaml에서 Torch 모듈을 오류 없이 임포트할 수 있어야 합니다.

선택 사항: 실험 추적 설정

머신러닝 실험을 추적할 계획이라면 Hugging Face와 함께 Weights & Biases를 설정할 수 있습니다:

# Install Weights & Biases
pip install wandb

# Log in to Weights & Biases (this will prompt for your API key)
wandb login

# Install and log in to Hugging Face
pip install huggingface_hub
huggingface-cli login

W&B API 키는 https://wandb.ai/settings 에서, Hugging Face 토큰은 https://huggingface.co/settings/tokens 에서 받을 수 있습니다.

예제 프로젝트: XOR 신경망

XOR 함수를 학습하도록 신경망을 훈련하는 Hello World OCaml 프로젝트를 만들어 봅시다. 먼저 새 프로젝트를 생성하십시오:

# Create project directory
mkdir xor_example
cd xor_example

# Initialize dune project
dune init project xor_example
cd xor_example

이제 새 파일 bin/main.ml을 생성하십시오:

open Base
open Torch

(* XOR truth table *)
let input_data =
  Tensor.of_float2 [|
    [|0.; 0.|];
    [|0.; 1.|];
    [|1.; 0.|];
    [|1.; 1.|];
  |]

let target_data =
  Tensor.of_float2 [|
    [|0.|];
    [|1.|];
    [|1.|];
    [|0.|];
  |]

let () =
  (* Create a simple network with one hidden layer *)
  let vs = Var_store.create ~name:"xor" () in
  let hidden = Layer.linear vs ~input_dim:2 4 ~activation:Relu in
  let output = Layer.linear vs ~input_dim:4 1 ~activation:Sigmoid in

  (* Define the model *)
  let model x =
    Layer.forward hidden x
    |> Layer.forward output
  in

  (* Training configuration *)
  let learning_rate = 0.05 in
  let optimizer = Optimizer.adam vs ~learning_rate in

  (* Training loop *)
  for epoch = 1 to 2000 do
    (* Forward pass *)
    let predicted = model input_data in
    let loss = Tensor.mse_loss predicted target_data in

    (* Backward pass *)
    Optimizer.backward_step optimizer ~loss;

    (* Print progress every 100 epochs *)
    if epoch % 100 = 0 then
      Stdio.printf "Epoch %d: Loss = %f\n%!"
        epoch (Tensor.float_value loss)
  done;

  (* Print final results *)
  let final_output = model input_data in
  Stdio.printf "\nFinal predictions:\n";
  Tensor.print final_output

bin 디렉터리의 dune 파일을 수정하십시오:

(executable
 (name main)
 (libraries base stdio torch))

프로젝트를 빌드하고 실행하십시오:

# Build the project
dune build

# Run the example
dune exec xor_example

정상적으로 실행된다면 축하합니다. 이제 OCaml로 직접 훈련 실행을 조율할 준비가 되었습니다.

Keyboard shortcuts

Stephen Diehl 블로그 한국어 번역

OCaml용 PyTorch 설정하기

DeepSpeed 설치

Megatron-LM 설치