GitHub - NousResearch/tinker-atropos: Standalone repo for our Atropos integration with Thinking Machines Tinker API (https://thinkingmachines.ai/tinker/)
핵심 포인트
- 1tinker-atropos는 NousResearch의 Atropos와 Thinking Machines의 Tinker API를 연결하는 통합 레이어로, 클라우드 기반 분산 RLHF 훈련을 가능하게 합니다.
- 2이 패키지는 사용자가 로컬 머신에서 Atropos 환경을 활용하여 모델 훈련을 원활하게 수행할 수 있도록 컴퓨팅 관리 및 인프라 복잡성을 추상화합니다.
- 3기존 Atropos 환경은 표준 구성에 Tinker 관련 설정을 추가하는 방식으로 쉽게 호환되며, 유연한 YAML 구성 파일을 통해 훈련 파라미터를 제어할 수 있습니다.
tinker-atropos는 NousResearch의 Atropos 프로젝트와 Thinking Machines의 Tinker API를 연결하는 통합 레이어(integration layer) 패키지입니다. 이 프로젝트의 핵심 목표는 사용자가 로컬 머신에서 Atropos 환경을 통해 모델 훈련을 원활하게 수행할 수 있도록 컴퓨팅 관리 및 인프라 관련 복잡성을 추상화하는 것입니다.
설치 및 빠른 시작:
패키지는 pip install -e . 명령어를 통해 설치할 수 있습니다. 훈련을 시작하기 위해서는 세 가지 터미널 창을 사용한 워크플로우가 필요합니다.
- 터미널 1:
AtroposAPI를 시작합니다 (run-api). - 터미널 2:
TinkerAPI 키를 환경 변수TINKER_API_KEY로 설정한 후,launch_training.py스크립트를 실행하여 훈련을 시작합니다. 이때configs/default.yaml과 같은 설정 파일을--config플래그로 지정할 수 있습니다. - 터미널 3:
Atropos환경(예:tinker_atropos/environments/gsm8k_tinker.py)을serve모드로 실행하고, 이때 훈련에서 사용한 것과 동일한 설정 파일을--config플래그로 지정합니다. 이 설정 파일은 환경 및 추론 서버 구성을 포함합니다.
예시로는 GSM8k 데이터셋에 대해 Llama-3.1-8B-Instruct 모델을 50단계(step)로 훈련하는 과정이 제공됩니다.
핵심 동작 방식 (How It Works):tinker-atropos의 핵심은 Atropos 환경이 --config 플래그를 통해 Tinker 설정(표준 Atropos 형식에 tinker 섹션이 추가된 형태)을 로드할 수 있도록 하는 것입니다. 이 설정 파일은 세 가지 주요 섹션으로 구성됩니다.
env:Atropos환경 자체의 구성(예:group_size,batch_size,tokenizer_name).openai: 추론 서버(inference server)의 구성(예:model_name,base_url).tinker:Tinker서비스와의 연동 및 훈련 관련 특정 설정(예:lora_rank,learning_rate,max_token_trainer_length,wandb_project등).
여기서 중요한 점은 tinker 섹션은 트레이너(trainer)에 의해 사용되며, Atropos 환경은 이 섹션을 무시한다는 것입니다. 따라서, managed_server를 사용하는 기존의 모든 Atropos 환경은 별도의 수정 없이 Tinker와 호환됩니다. 환경은 managed_server를 사용할 때 self.tokenizer.eos_token_id와 같은 stop 시퀀스를 반드시 포함해야 합니다.
async with self.server.managed_server(tokenizer=self.tokenizer) as managed:
chat_completion = await managed.chat_completion(
messages=messages,
n=self.config.group_size,
max_tokens=self.config.max_token_length,
temperature=1.0,
stop=[self.tokenizer.eos_token_id], # Add stop sequences
)설정 및 사용법:
트레이너와 환경 모두 YAML 설정 파일을 사용하여 파라미터 일관성을 보장합니다. CLI를 통해 --num-steps 100, --no-wandb와 같이 개별 파라미터를 오버라이드(override)할 수도 있습니다.TinkerAtroposConfig 클래스를 사용하여 YAML 파일로부터 설정을 로드하거나 기본값을 사용할 수 있으며, TinkerAtroposTrainer를 초기화하고 await trainer.run()을 통해 훈련을 실행하는 프로그래밍 방식의 사용도 지원합니다.
훈련 완료 시, 최종 가중치는 형식의 Tinker 경로로 출력되며, 이는 tinker_atropos/utils/download_weights.py 스크립트를 통해 다운로드할 수 있습니다.
tinker-atropos는 Atropos의 강화 학습(RL) 환경과 Tinker의 클라우드 기반 컴퓨팅 인프라를 연결하여, 사용자가 로컬에서 개발하고 클라우드에서 스케일업(scale up)할 수 있는 유연한 ML 훈련 워크플로우를 제공합니다.