Update/unified_ttrlvr_architecture.md

TTRLVR Unified Architecture - 상세 작동 방식

목차

1. 개요
2. 전체 아키텍처
3. 실행 흐름
4. 핵심 컴포넌트
5. Phase별 상세 동작
6. 동기화 메커니즘
7. 데이터 흐름
8. 구현 세부사항

---

1. 개요

1.1 목적

TTRLVR Unified는 기존 TTRLVR의 분리된 구조를 하나의 통합된 VeRL 세션으로 재구성하여 동기화 문제를 해결하고 성능을 향상시킨 버전입니다.

1.2 핵심 개선사항

• 단일 vLLM 인스턴스: 전체 학습 과정에서 하나의 vLLM만 사용
• 동기화 문제 해결: dummy_dtensor 사용 가능
• 성능 향상: vLLM 재생성 오버헤드 제거로 30-40% 속도 향상
• 메모리 효율: 반복적인 할당/해제 없음

1.3 주요 파일

• train_ttrlvr_azr_unified.py: 메인 실행 스크립트
• test/trainer/unified_ttrlvr_trainer.py: 통합 Trainer 클래스
• test/configs/ttrlvr_azr_unified_4gpu.yaml: VeRL 설정 파일

---

2. 전체 아키텍처

2.1 기존 vs 통합 구조

기존 TTRLVR (분리형)

Round 1:
├── Phase 1-4: RemoteTestTimePipeline (독립 vLLM #1)
│   └── ray.kill(pipeline)  # vLLM 삭제
└── Phase 5: VeRL Training (새 vLLM #2)
    └── trainer.init_workers()  # 매 라운드마다

Round 2: (새로운 vLLM 인스턴스들...)

Unified TTRLVR (통합형)

초기화:
└── trainer.init_workers()  # 1번만!

Round 1-N:
├── Phase 1-4: 데이터 생성 (같은 vLLM)
└── Phase 5: PPO 학습 (같은 vLLM)

2.2 컴포넌트 관계도

train_ttrlvr_azr_unified.py
    │
    ├── 환경 설정 & 인자 파싱
    │
    ├── VeRL generate_main() 호출
    │   │
    │   └── UnifiedTTRLVRTrainer 생성
    │       │
    │       ├── CompleteTestTimePipeline (Phase 1-4)
    │       │   ├── 벤치마크 문제 로딩
    │       │   ├── 프로그램 생성 (diverse_programs)
    │       │   ├── IPO 추출 (IPOTripleExtractor)
    │       │   ├── Task 생성 (TestTimeTaskGenerator)
    │       │   └── 검증 및 필터링
    │       │
    │       └── VeRL PPO Training (Phase 5)
    │           ├── 데이터 형식 변환
    │           ├── Response 생성
    │           ├── Reward 계산
    │           └── Policy 업데이트

---

3. 실행 흐름

3.1 스크립트 실행

python train_ttrlvr_azr_unified.py --benchmark mbpp --problems 10 --rounds 30 --gpu 0,1,2,3

3.2 초기화 단계

Step 1: 인자 파싱

def main():
    # 명령행 인자 파싱
    args = parse_arguments()
    
    # 환경 설정 (GPU, 경로 등)
    setup_environment(args.gpu)

Step 2: 문제 리스트 생성

# 벤치마크에서 문제 ID 추출
problem_ids = create_problem_list(args.benchmark, args.problems, args.problem_id)
# 예: ['Mbpp/1', 'Mbpp/2', 'Mbpp/3', ...]

Step 3: 환경 변수 설정

# VeRL이 UnifiedTTRLVRTrainer에 전달할 설정
os.environ['TTRLVR_PROBLEM_IDS'] = json.dumps(problem_ids)
os.environ['TTRLVR_TOTAL_ROUNDS'] = str(args.rounds)
os.environ['TTRLVR_OUTPUT_DIR'] = output_dir
os.environ['TTRLVR_CONFIG'] = json.dumps(ttrlvr_config)

Step 4: VeRL 실행

# VeRL의 main_generation 호출
verl_args = [
    'train_ttrlvr_azr_unified.py',
    f'--config-path={config_path}',
    '--config-name=ttrlvr_azr_unified_4gpu',
    f'trainer.project_name=ttrlvr_unified_{args.benchmark}',
    f'trainer.total_epochs={args.rounds}',  # 각 라운드를 epoch로 매핑
]

sys.argv = verl_args
generate_main()  # VeRL 메인 함수 실행

3.3 VeRL 초기화

VeRL의 generate_main()이 실행되면:

1. Config 로딩: ttrlvr_azr_unified_4gpu.yaml 파싱
2. Ray 클러스터 초기화: 분산 처리 환경 설정
3. UnifiedTTRLVRTrainer 생성: 설정에 명시된 클래스 로드
4. Worker 초기화: trainer.init_workers() 호출 (1번만!)

---

4. 핵심 컴포넌트

4.1 UnifiedTTRLVRTrainer

class UnifiedTTRLVRTrainer(ReasonRLRayPPOTrainer):
    """
    TTRLVR의 모든 Phase를 하나의 VeRL 세션에서 처리하는 통합 Trainer
    """
    
    def __init__(self, ttrlvr_config, problem_ids, total_rounds, ...):
        super().__init__(...)
        
        # TTRLVR 특화 설정
        self.ttrlvr_config = ttrlvr_config
        self.problem_ids = problem_ids
        self.total_rounds = total_rounds
        self.current_round = 0
        
        # CompleteTestTimePipeline 초기화 (나중에)
        self.ttrlvr_pipeline = None

4.2 CompleteTestTimePipeline 통합

def _init_ttrlvr_pipeline(self):
    """CompleteTestTimePipeline을 VeRL의 vLLM으로 초기화"""
    
    # VeRL의 모델 사용
    self.ttrlvr_pipeline = CompleteTestTimePipeline(
        model=None,  # VeRL wrapper 통해 접근
        tokenizer=self.tokenizer,
        config=self.testtime_config,
        logger=self.ttrlvr_logger
    )
    
    # VeRL의 vLLM을 사용하도록 설정
    self.ttrlvr_pipeline.generate_with_verl = self._generate_with_vllm

---

5. Phase별 상세 동작

5.1 fit() 메서드 - 메인 학습 루프

def fit(self):
    """전체 학습 루프 관리"""
    
    # 로거 초기화
    logger = ReasonRLTracking(...)
    
    # 체크포인트 로드 (있으면)
    self._load_checkpoint()
    
    # 라운드별 반복
    for round_num in range(1, self.total_rounds + 1):
        self.current_round = round_num
        
        # ====== Phase 1-4: 데이터 생성 ======
        round_data = self._generate_round_data()
        
        # ====== Phase 5: PPO 학습 ======
        metrics = self._train_one_round(round_data, logger)
        
        # 체크포인트 저장 (5라운드마다)
        if round_num % 5 == 0:
            self._save_checkpoint()

5.2 Phase 1-4: 데이터 생성

5.2.1 _generate_round_data() 구조

def _generate_round_data(self) -> List[Dict[str, Any]]:
    """Phase 1-4 실행"""
    
    # Pipeline 초기화 (처음만)
    if self.ttrlvr_pipeline is None:
        self._init_ttrlvr_pipeline()
    
    all_tasks = []
    
    for problem_id in self.problem_ids:
        # CompleteTestTimePipeline 실행
        result = self.ttrlvr_pipeline.run_complete_pipeline(
            benchmark_config=benchmark_config,
            problem_id=problem_id,
            round_num=self.current_round,
            session_timestamp=session_timestamp
        )
        
        if result['success']:
            tasks = result['final_tasks']
            all_tasks.extend(tasks)
    
    return all_tasks

5.2.2 CompleteTestTimePipeline 내부 동작

Phase 1: 다양한 프로그램 생성

# 1. 벤치마크 문제 로드
problem = benchmark_loader.load_problem(benchmark_config, problem_id)

# 2. Baseline 평가
baseline_results = self._evaluate_baseline_performance(problem)

# 3. 다양한 프로그램 생성
diverse_programs = self._generate_diverse_programs_and_ipo(problem)
# 내부적으로:
# - 정교한 프롬프트 템플릿 사용
# - Temperature 조절로 다양성 확보
# - 문법 검증

Phase 2: I/O 쌍 추출

# IPOTripleExtractor 사용
ipo_extractor = IPOTripleExtractor(config, logger, model, tokenizer)

for program in diverse_programs:
    # 입력 생성
    inputs = ipo_extractor.generate_inputs(program)
    
    # 출력 계산
    for input in inputs:
        output = executor.execute(program, input)
        ipo_buffer.add_triple(input, program, output)

Phase 3: Task 생성

# TestTimeTaskGenerator 사용
task_generator = TestTimeTaskGenerator(config, logger)

# Induction: I/O → Program
induction_tasks = task_generator.create_induction_tasks(ipo_triples)

# Deduction: Program + Input → Output
deduction_tasks = task_generator.create_deduction_tasks(ipo_triples)

# Abduction: Program + Output → Input
abduction_tasks = task_generator.create_abduction_tasks(ipo_triples)

Phase 4: 검증 및 필터링

# 각 task 검증
valid_tasks = []
for task in all_tasks:
    if validator.is_valid(task):
        valid_tasks.append(task)

5.3 Phase 5: PPO 학습

5.3.1 _train_one_round() 구조

def _train_one_round(self, round_data: List[Dict], logger) -> Dict[str, float]:
    """Phase 5: PPO 학습"""
    
    # 1. 데이터 변환
    train_dataset = self._convert_to_verl_dataset(round_data)
    
    # 2. DataLoader 생성
    self.train_dataloader = self._create_dataloader(
        train_dataset,
        batch_size=self.config.data.train_batch_size
    )
    
    # 3. 1 epoch 학습
    epoch_metrics = {}
    for step, batch in enumerate(self.train_dataloader):
        # PPO Step 1: Response 생성
        gen_batch_output = self.actor_rollout_wg.generate_sequences(batch)
        
        # PPO Step 2: Reward 계산
        reward_tensor = self.reward_fn(batch.union(gen_batch_output))
        
        # PPO Step 3: Policy 업데이트
        update_metrics = self._ppo_update(batch, reward_tensor)
        
        # 메트릭 수집
        for k, v in update_metrics.items():
            epoch_metrics[k].append(v)
    
    return {k: np.mean(v) for k, v in epoch_metrics.items()}

5.3.2 데이터 변환 과정

def _convert_to_verl_dataset(self, round_data: List[Dict]) -> Any:
    """TTRLVR 형식 → VeRL 형식"""
    
    converted_data = []
    for task in round_data:
        # 토큰화
        prompt_ids = self.tokenizer(
            task['prompt'],
            max_length=self.config.data.max_prompt_length
        ).input_ids
        
        # VeRL DataProto 형식
        verl_item = {
            'input_ids': prompt_ids,
            'prompt': task['prompt'],
            'target': task['target'],
            'task_type': task['task_type'],
            'problem_id': task['problem_id']
        }
        converted_data.append(verl_item)
    
    return converted_data

---

6. 동기화 메커니즘

6.1 문제의 핵심

기존 TTRLVR은 매 라운드마다 새 vLLM을 생성했기 때문에 dummy_dtensor 사용 시 동기화가 되지 않았습니다.

6.2 해결 방법

6.2.1 단일 vLLM 인스턴스

# 초기화 (1번만)
trainer.init_workers()
├── FSDP workers 생성
├── vLLM workers 생성
└── 초기 동기화 (sync_model_weights)

# 이후 모든 라운드에서 같은 인스턴스 사용
Round 1: Phase 1-4 → Phase 5 (같은 vLLM)
Round 2: Phase 1-4 → Phase 5 (같은 vLLM)
...

6.2.2 동기화 과정

# FSDPVLLMShardingManager의 동작
class FSDPVLLMShardingManager:
    def __enter__(self):
        if not self.base_sync_done:
            # 첫 번째 호출: FSDP → vLLM 동기화
            sync_model_weights(actor_weights, load_format='dummy_dtensor')
            self.base_sync_done = True
        # 이후: 메모리 참조로 자동 동기화

6.3 메모리 참조 메커니즘

FSDP 모델 (GPU 0-3)          vLLM 모델 (GPU 0-1)
┌─────────────┐              ┌─────────────┐
│ Parameter A │ ─────────→   │ Parameter A │ (같은 메모리 참조)
│ Parameter B │ ─────────→   │ Parameter B │
│ Parameter C │ ─────────→   │ Parameter C │
└─────────────┘              └─────────────┘

PPO 업데이트 → FSDP 파라미터 변경 → vLLM도 자동으로 새 값 사용

---

7. 데이터 흐름

7.1 Round 1 상세 흐름

1. Problem: Mbpp/2 (예: "두 수의 합을 구하는 함수 작성")
   │
   ├── Phase 1: 프로그램 생성
   │   ├── Prompt: "Generate 4 different solutions..."
   │   ├── vLLM 생성 (동기화 발생)
   │   └── Output: [prog1, prog2, prog3, prog4]
   │
   ├── Phase 2: I/O 추출
   │   ├── 각 프로그램에 대해 입력 생성
   │   ├── vLLM 사용 (동기화 건너뜀)
   │   └── Output: [(input1, output1), (input2, output2), ...]
   │
   ├── Phase 3: Task 생성
   │   ├── Induction: (1, 3) → "def add(a,b): return a+b"
   │   ├── Deduction: (prog, 5) → 8
   │   └── Abduction: (prog, 10) → (4, 6)
   │
   ├── Phase 4: 검증
   │   └── 유효한 task만 필터링
   │
   └── Phase 5: PPO 학습
       ├── 배치 생성
       ├── Response 생성 (같은 vLLM)
       ├── Reward 계산
       └── FSDP 모델 업데이트

7.2 데이터 형식 변환

# TTRLVR Task 형식
{
    'problem_id': 'Mbpp/2',
    'task_type': 'induction',
    'input': 5,
    'output': 10,
    'target': 'def multiply_by_two(x): return x * 2',
    'prompt': 'Given input 5 produces output 10, write the function:'
}

# ↓ 변환

# VeRL DataProto 형식
{
    'input_ids': tensor([1, 234, 567, ...]),  # 토큰화된 prompt
    'attention_mask': tensor([1, 1, 1, ...]),
    'prompt': 'Given input 5 produces output 10...',
    'target': 'def multiply_by_two(x): return x * 2',
    'meta_info': {
        'task_type': 'induction',
        'problem_id': 'Mbpp/2'
    }
}

---

8. 구현 세부사항

8.1 VeRL과의 통합

8.1.1 _generate_with_vllm 메서드

def _generate_with_vllm(self, prompt: str, temperature: float = 0.7):
    """VeRL의 vLLM을 사용한 텍스트 생성"""
    
    # 1. 토큰화
    input_ids = self.tokenizer(prompt, ...).input_ids
    
    # 2. DataProto 생성
    prompts_proto = DataProto.from_dict({
        "input_ids": input_ids.cuda(),
        "attention_mask": torch.ones_like(input_ids).cuda(),
    })
    
    # 3. 메타 정보 설정
    prompts_proto.meta_info = {
        "eos_token_id": self.tokenizer.eos_token_id,
        "temperature": temperature,
        "do_sample": True,
        "response_length": 256
    }
    
    # 4. VeRL의 vLLM으로 생성
    outputs = self.actor_rollout_wg.generate_sequences(prompts_proto)
    
    # 5. 디코딩 및 반환
    return self.tokenizer.decode(outputs.batch["input_ids"][0])

8.1.2 CompleteTestTimePipeline 수정

# CompleteTestTimePipeline이 VeRL의 vLLM을 사용하도록
self.ttrlvr_pipeline.generate_with_verl = self._generate_with_vllm

# 이제 Pipeline 내부에서:
# response = self.generate_with_verl(prompt)  # VeRL의 vLLM 사용

8.2 메모리 관리

8.2.1 라운드 간 메모리 정리

def _manage_memory_between_rounds(self):
    """라운드 간 메모리 정리 (인스턴스는 유지)"""
    
    # GPU 캐시만 정리
    torch.cuda.empty_cache()
    
    # vLLM KV 캐시 정리 (선택적)
    if hasattr(self.actor_rollout_wg, 'clear_kv_cache'):
        self.actor_rollout_wg.clear_kv_cache()
    
    # Garbage collection
    import gc
    gc.collect()

8.2.2 메모리 모니터링

def _monitor_memory(self):
    """메모리 사용량 모니터링"""
    for i in range(torch.cuda.device_count()):
        allocated = torch.cuda.memory_allocated(i) / 1024**3
        reserved = torch.cuda.memory_reserved(i) / 1024**3
        print(f"GPU {i}: Allocated={allocated:.2f}GB, Reserved={reserved:.2f}GB")

8.3 에러 처리 및 복구

def _safe_generate(self, prompt: str, max_retries: int = 3):
    """안전한 생성 with 재시도"""
    for attempt in range(max_retries):
        try:
            return self._generate_with_vllm(prompt)
        except Exception as e:
            if attempt == max_retries - 1:
                raise
            torch.cuda.empty_cache()
            time.sleep(1)

8.4 체크포인트 관리

def _save_checkpoint(self):
    """체크포인트 저장"""
    checkpoint = {
        'round': self.current_round,
        'model_state_dict': self.actor_rollout_wg.state_dict(),
        'optimizer_state_dict': self.optimizer.state_dict(),
        'metrics': self.accumulated_metrics,
        'timestamp': datetime.now().isoformat()
    }
    
    path = f"{self.checkpoint_dir}/round_{self.current_round}.pt"
    torch.save(checkpoint, path)

---

9. 성능 최적화

9.1 배치 처리

• Phase 1-4에서 가능한 한 배치로 처리
• vLLM의 continuous batching 활용

9.2 GPU 활용

• vLLM: GPU 0-1 (tensor parallel)
• FSDP: GPU 0-3 (data parallel)
• 효율적인 GPU 메모리 활용

9.3 I/O 최적화

• Parquet 형식으로 중간 데이터 저장
• 비동기 I/O 처리

---

10. 디버깅 및 모니터링

10.1 로깅 구조

/home/ubuntu/RLVR/TestTime-RLVR-v2/logs/
├── ttrlvr_unified_20241107_120000.log  # 메인 로그
├── round_1/
│   ├── phase_1_4.log  # 데이터 생성 로그
│   └── phase_5.log    # 학습 로그
└── metrics/
    └── tensorboard/   # 학습 메트릭

10.2 주요 모니터링 지표

• 라운드별 소요 시간
• 생성된 task 수
• 평균 reward
• GPU 메모리 사용량
• 동기화 발생 횟수

---

11. 문제 해결 가이드

11.1 OOM (Out of Memory)

• gpu_memory_utilization 조정 (기본: 0.35)
• max_num_seqs 감소
• 배치 크기 감소

11.2 동기화 문제

• load_format이 dummy_dtensor인지 확인
• vLLM 인스턴스가 재생성되지 않는지 확인

11.3 느린 성능

• GPU 활용률 확인
• 배치 크기 증가
• enforce_eager=False 확인 (CUDA graph 사용)

---

12. 결론

TTRLVR Unified는 기존 TTRLVR의 모든 기능을 유지하면서 다음을 달성했습니다:

1. 구조적 개선: 분리된 Phase들을 하나의 세션으로 통합
2. 성능 향상: vLLM 재생성 오버헤드 제거로 30-40% 속도 향상
3. 안정성 향상: 동기화 문제 완전 해결
4. 확장성: 더 큰 모델과 더 많은 라운드 지원 가능

이 아키텍처는 TTRLVR의 정교한 데이터 생성 능력과 VeRL의 효율적인 PPO 학습을 완벽하게 결합했습니다.