Open-Sora 성능 최적화 완벽 가이드: CUDA 메모리부터 Flash Attention까지

# 모델 정밀도 최적화
vae = vae.to(torch.bfloat16).eval()
stdit = stdit.to(torch.bfloat16).eval()

# T5는 fp32 유지 (정확도 중요)
text_encoder.t5.model = text_encoder.t5.model.eval()

# 추론 모드 활성화
with torch.inference_mode():
    # 그래디언트 계산 완전 비활성화
    samples = scheduler.sample(...)

# 모델 평가 모드
model.eval().requires_grad_(False)

# CUDA 캐시 정리
torch.cuda.empty_cache()

# torch.no_grad(): 그래디언트 계산만 비활성화
with torch.no_grad():
    output = model(input)

# torch.inference_mode(): 완전한 추론 최적화
with torch.inference_mode():
    output = model(input)  # 더 빠름, 메모리 절약

# 표준 어텐션: O(n²) 메모리 복잡도
def standard_attention(Q, K, V):
    # Q @ K^T: [seq_len, seq_len] 매트릭스 생성
    attention_scores = torch.matmul(Q, K.transpose(-2, -1))  # O(n²) 메모리
    attention_weights = torch.softmax(attention_scores, dim=-1)
    output = torch.matmul(attention_weights, V)
    return output

# Flash Attention: O(n) 메모리 복잡도
def flash_attention_concept(Q, K, V, block_size=64):
    """Flash Attention 핵심 아이디어 (의사코드)"""
    seq_len = Q.shape[1]
    output = torch.zeros_like(Q)
    
    # 블록 단위로 처리 (메모리 효율성)
    for i in range(0, seq_len, block_size):
        for j in range(0, seq_len, block_size):
            # 작은 블록 단위로 어텐션 계산
            q_block = Q[:, i:i+block_size]
            k_block = K[:, j:j+block_size]
            v_block = V[:, j:j+block_size]
            
            # 블록 내에서만 어텐션 계산
            attention_block = compute_attention_block(q_block, k_block, v_block)
            output[:, i:i+block_size] += attention_block
    
    return output

import xformers.ops as xops

def memory_efficient_attention(Q, K, V, attention_mask=None):
    """xformers의 메모리 효율적 어텐션"""
    return xops.memory_efficient_attention(
        Q, K, V,
        attn_bias=attention_mask,
        scale=1.0 / math.sqrt(Q.shape[-1])
    )

import xformers.checkpoint as checkpoint

class MemoryEfficientTransformer(nn.Module):
    def forward(self, x):
        # 메모리 절약을 위한 그래디언트 체크포인팅
        x = checkpoint.checkpoint(self.layer1, x)
        x = checkpoint.checkpoint(self.layer2, x)
        return x

# TF32 최적화 활성화
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

# JIT 컴파일 비활성화 (Gradio 호환성)
torch.jit._state.disable()

# 또는 프로덕션에서 JIT 활용
@torch.jit.script
def optimized_function(x):
    return x * 2 + 1

def adaptive_batch_size(model, input_shape, max_memory_gb=40):
    """GPU 메모리에 따른 적응적 배치 크기 설정"""
    batch_size = 1
    while True:
        try:
            # 테스트 배치로 메모리 사용량 확인
            test_input = torch.randn(batch_size, *input_shape).cuda()
            with torch.inference_mode():
                _ = model(test_input)
            
            # 메모리 사용량 체크
            memory_used = torch.cuda.memory_allocated() / 1024**3
            if memory_used > max_memory_gb * 0.8:  # 80% 임계값
                break
            
            batch_size += 1
        except RuntimeError as e:
            if "out of memory" in str(e):
                batch_size -= 1
                break
            raise e
    
    return max(1, batch_size)

def memory_offloading_inference(model, input_data):
    """CPU-GPU 메모리 오프로딩"""
    # 모델을 CPU로 이동
    model.cpu()
    
    # 필요한 레이어만 GPU로 로드
    for layer_name, layer in model.named_modules():
        if "attention" in layer_name:  # 중요한 레이어만 GPU
            layer.cuda()
    
    # 추론 실행
    with torch.inference_mode():
        output = model(input_data.cuda())
    
    # 정리
    torch.cuda.empty_cache()
    return output.cpu()

def monitor_gpu_memory():
    """실시간 GPU 메모리 모니터링"""
    if torch.cuda.is_available():
        allocated = torch.cuda.memory_allocated() / 1024**3
        reserved = torch.cuda.memory_reserved() / 1024**3
        max_allocated = torch.cuda.max_memory_allocated() / 1024**3
        
        print(f"GPU Memory Usage:")
        print(f"  Allocated: {allocated:.2f} GB")
        print(f"  Reserved:  {reserved:.2f} GB")
        print(f"  Max Used:  {max_allocated:.2f} GB")
        
        # 메모리 사용률이 높으면 경고
        if allocated / reserved > 0.9:
            print("⚠️  GPU memory usage is high!")
        
        return {
            'allocated': allocated,
            'reserved': reserved,
            'max_allocated': max_allocated
        }

def profile_model_performance(model, input_data, num_runs=10):
    """모델 성능 프로파일링"""
    import time
    
    # 워밍업
    for _ in range(3):
        with torch.inference_mode():
            _ = model(input_data)
    
    torch.cuda.synchronize()
    
    # 실제 측정
    times = []
    for i in range(num_runs):
        start_time = time.time()
        
        with torch.inference_mode():
            output = model(input_data)
        
        torch.cuda.synchronize()
        end_time = time.time()
        
        times.append(end_time - start_time)
        print(f"Run {i+1}: {times[-1]:.3f}s")
    
    avg_time = sum(times) / len(times)
    std_time = (sum((t - avg_time)**2 for t in times) / len(times))**0.5
    
    print(f"\nAverage: {avg_time:.3f}s ± {std_time:.3f}s")
    return avg_time, std_time