[ultralytics] SAM3 텍스트 프롬프팅 VRAM 최적화: 청크 기반 처리로 메모리 사용량 대폭 감소

PR 링크: ultralytics/ultralytics#26281 상태: Merged | 변경: +20 / -0

들어가며

최근 컴퓨터 비전 분야에서 텍스트 프롬프트를 활용한 이미지 이해 및 생성 모델의 중요성이 날로 커지고 있습니다. 특히 Ultralytics의 SAM3 모델은 텍스트 설명을 기반으로 객체를 탐지하고 분할하는 강력한 기능을 제공합니다. 하지만 대규모 텍스트 프롬프트 세트를 처리할 때 발생하는 높은 VRAM 사용량은 특히 중급 이하의 GPU 환경에서 심각한 병목 현상을 야기했습니다. 이전 SAM3 구현에서는 모든 텍스트 프롬프트를 단일 배치로 처리하여, 프롬프트 수가 증가함에 따라 VRAM 사용량이 선형적으로 증가하는 문제가 있었습니다. 예를 들어, 80개의 텍스트 프롬프트는 약 13GiB, 128개는 약 20GiB의 VRAM을 요구하여 Out-of-Memory (OOM) 오류를 빈번하게 발생시켰습니다.

이번 PR은 이러한 문제를 해결하기 위해 텍스트 프롬프트를 작은 '청크(chunk)' 단위로 나누어 처리하는 혁신적인 방식을 도입했습니다. 이를 통해 SAM3 모델의 텍스트 프롬프팅 기능 사용 시 피크 VRAM 사용량을 획기적으로 줄이고, 더 넓은 범위의 하드웨어에서 모델을 효율적으로 사용할 수 있게 되었습니다.

코드 분석: 청크 기반 텍스트 프롬프트 처리

이번 최적화의 핵심은 SAM3SemanticModel.forward_grounding 메서드 내에서 텍스트 프롬프트 처리 방식을 변경한 것입니다. 이전에는 모든 텍스트 프롬프트가 하나의 배치로 처리되었지만, 이제는 SAM3SemanticModel.max_text_batch (기본값 16) 설정값에 따라 프롬프트들을 그룹화하여 순차적으로 처리합니다.

ultralytics/models/sam/sam3/sam3_image.py

가장 핵심적인 변경은 forward_grounding 메서드에 도입된 청크 처리 로직입니다. 텍스트 프롬프트의 수가 max_text_batch를 초과하고, geometric_prompt (박스 프롬프트 등)가 없거나 비어있는 경우에만 이 청크 처리 로직이 활성화됩니다.

Before:

 def forward_grounding( self, backbone_out: dict[str, torch.Tensor], text_ids: torch.Tensor, geometric_prompt: Prompt = None ): """Forward pass for grounding (detection + segmentation) given input images and text."""
 # backbone features are expanded to one batch item per prompt, and downstream encoder attention / pixel-decoder convs / mask einsum materialize per-prompt tensors.
 # Peak VRAM grows linearly ~145 MB per prompt (FP16): 13 GiB at 80 prompts, ~20 GiB at 128 -> OOM on mid-range GPUs.
 # ... (기존 코드) ...

After:

 def forward_grounding( self, backbone_out: dict[str, torch.Tensor], text_ids: torch.Tensor, geometric_prompt: Prompt = None ): """Forward pass for grounding (detection + segmentation) given input images and text."""
+ if len(text_ids) > self.max_text_batch and (
+ geometric_prompt is None or geometric_prompt.box_embeddings.shape[0] == 0
+ ):
+ # Chunk text prompts to bound peak VRAM; prompts are independent batch items so results concatenate
+ # exactly. An empty geometric prompt carries no boxes, so slicing its batch dim preserves numerics.
+ outs = []
+ for i in range(0, len(text_ids), self.max_text_batch):
+ prompt = None
+ if geometric_prompt is not None:
+ prompt = Prompt(
+ box_embeddings=geometric_prompt.box_embeddings[:, i : i + self.max_text_batch],
+ box_mask=geometric_prompt.box_mask[i : i + self.max_text_batch],
+ box_labels=geometric_prompt.box_labels[:, i : i + self.max_text_batch],
+ )
+ outs.append(self.forward_grounding(backbone_out, text_ids[i : i + self.max_text_batch], prompt))
+ return {
+ k: torch.cat([o[k] for o in outs]) if isinstance(v, torch.Tensor) else v for k, v in outs[0].items()
+ }
 # backbone_out, img_ feats, img_ pos_embeds, vis_feat_sizes = SAM2Model._prepare_backbone_features(
 # self, backbone_out, batch=len(text_ids)
 # )

새로운 로직은 다음과 같이 작동합니다:

1. 조건 확인: 처리해야 할 text_ids의 수가 self.max_text_batch보다 많고, geometric_prompt가 없거나 비어있는 경우에만 청크 처리를 수행합니다. 이는 텍스트 프롬프트만으로 구성된 경우 VRAM 부담이 크기 때문입니다. 박스 프롬프트 등이 함께 사용될 경우, 이는 일반적으로 VRAM 사용량이 상대적으로 적거나, 박스 프롬프트 자체가 배치 크기를 제어하는 역할을 할 수 있어 별도의 청크 처리가 필요 없을 수 있습니다.
2. 청크 반복: range(0, len(text_ids), self.max_text_batch)를 사용하여 text_ids를 max_text_batch 크기의 청크로 나눕니다.
3. 재귀 호출: 각 청크에 대해 forward_grounding 메서드를 재귀적으로 호출합니다. 이때, geometric_prompt가 존재하면 해당 청크에 맞는 부분만 슬라이싱하여 전달합니다. geometric_prompt가 None인 경우, None으로 전달되어 해당 청크에 대한 처리가 진행됩니다.
4. 결과 취합: 각 청크에서 반환된 결과(outs)들을 torch.cat을 사용하여 프롬프트 차원(prompt dim)을 따라 연결합니다. 텍스트 프롬프트는 독립적인 배치 아이템이므로, 각 청크의 결과는 단순히 이어 붙여도 정확한 결과를 보장합니다.

또한, 비디오 프롬프팅(SAM3VideoSemanticPredictor)의 경우에도 동일한 로직이 적용됩니다. 비디오 예측기는 geometric_prompt=None 대신 빈 Prompt 객체를 전달하는데, 이 경우에도 새로운 청크 로직은 geometric_prompt.box_embeddings.shape[0] == 0 조건을 만족하여 정상적으로 작동합니다. 이는 geometric_prompt의 배치 차원을 슬라이싱하여 수치적 안정성을 유지하기 위함입니다. (더미 Prompt는 프롬프트별 CLS 토큰을 인코딩하므로 None과 직접 호환되지 않습니다.)

박스 프롬프트와 같이 실제 기하학적 정보가 포함된 프롬프트에 대한 호출은 변경되지 않고 그대로 유지됩니다. 이는 이러한 종류의 프롬프트는 이미 VRAM 사용량 측면에서 상대적으로 효율적이거나, 다른 제약 조건에 의해 처리 방식이 달라질 수 있기 때문입니다.

모델별로 이 max_text_batch 값을 튜닝할 수 있습니다. 예를 들어, predictor.model.max_text_batch = 8과 같이 설정하여 더 작은 청크로 처리할 수도 있습니다.

리뷰 피드백 반영

리뷰 과정에서 backbone_out의 배치 크기 관리에 대한 우려가 제기되었습니다. 특히 재귀 호출 시 _prepare_backbone_features 함수가 각 청크의 배치 크기에 맞게 backbone_fpn 등의 텐서를 확장하는데, 이 과정에서 발생하는 backbone_out의 배치 크기가 후속 처리(예: 비디오 트래커)에 영향을 미칠 수 있다는 지적이었습니다.

이에 대해 개발자는 _prepare_backbone_features 함수가 최상위 backbone_fpn만 확장하며, 실제 비디오 트래커가 사용하는 sam2_backbone_out (내부의 배치 크기 1인 딕셔너리)은 복사본으로 유지되므로 문제가 없다고 설명했습니다. 또한, 실제 런타임 체크 결과와 비디오 예측의 정확도를 근거로, 트래커의 self.tracker.backbone_out이 올바르게 단일 이미지 형태로 유지됨을 확인했습니다. 비록 반환되는 backbone_out의 최상위 키들이 청크 배치 크기(16)를 가질 수 있지만, 이미지 경로와 비디오 경로는 이들을 직접 사용하지 않으므로 문제가 없다는 결론입니다. 이러한 상세한 분석과 검증은 코드의 안정성을 높이는 데 기여했습니다.

왜 이게 좋은가: VRAM 최적화와 성능 향상

이 PR의 가장 큰 장점은 피크 VRAM 사용량을 획기적으로 감소시켰다는 점입니다. PR 설명에 제시된 증거를 보면 그 효과가 명확합니다.

이미지 경로 (RTX 6000D, FP16):

#text prompts	before	after
1	1.91 GiB	1.91 GiB
16	4.03 GiB	4.03 GiB
80	13.04 GiB	4.85 GiB
128	19.79 GiB	5.47 GiB

보시다시피, 프롬프트 수가 max_text_batch (16) 이하일 때는 거의 오버헤드가 없지만, 80개 또는 128개와 같이 많은 프롬프트를 사용할 경우 VRAM 사용량이 1/3 수준으로 감소했습니다. 특히 80개 프롬프트의 경우 13.04 GiB에서 4.85 GiB로, 128개 프롬프트의 경우 19.79 GiB에서 5.47 GiB로 줄어들어, 기존에 OOM을 발생시키던 환경에서도 충분히 실행 가능하게 되었습니다.

비디오 경로 (80개 프롬프트, 8프레임 비디오, imgsz=640):

before	after
13.75 GiB	6.14 GiB

비디오 처리에서도 유사한 VRAM 감소 효과를 보였습니다. 13.75 GiB에서 6.14 GiB로 줄어들어, 비디오 기반의 복잡한 시맨틱 분석 작업에서도 메모리 제약이 크게 완화되었습니다.

또한, 성능 저하 없이 오히려 약간의 속도 향상까지 관찰되었습니다. 80개 프롬프트에 대한 이미지 처리 시 지연 시간(latency)이 380ms에서 355ms로 감소했습니다. 이는 청크 처리 과정에서 발생하는 오버헤드가 전체 연산 시간을 단축시키는 데 기여했거나, 배치 크기 감소로 인한 GPU 캐시 효율성 증가 등의 부수적인 효과일 수 있습니다.

출력 결과의 동일성 또한 중요한 장점입니다. 청크 처리 방식은 결과의 정확성을 유지하면서 VRAM을 절약하는 것을 목표로 합니다. FP32, FP16 환경 모두에서 이전 방식과 비교했을 때 최대 오차(max box diff 1.2e-4 px, conf diff 3e-7)가 매우 작으며, 마스크 IoU도 0.9999 이상으로 거의 동일한 결과를 보였습니다. 이는 논리적 오류가 아닌 부동 소수점 연산 커널 선택에 따른 미세한 차이로 간주될 수 있습니다.

일반적 교훈

이 PR은 다음과 같은 일반적인 최적화 교훈을 제공합니다:

1. 배치 처리의 함정 인식: 모든 데이터를 한 번에 배치로 처리하는 것이 항상 최선은 아닙니다. 특히 메모리 집약적인 연산의 경우, 배치 크기가 선형적으로 VRAM 사용량에 영향을 미칠 수 있습니다. 이럴 때는 데이터를 작은 단위로 나누어 순차 처리하는 '청킹(chunking)' 기법이 효과적입니다.
2. 독립적인 작업 단위 활용: 텍스트 프롬프트와 같이 서로 독립적인 작업 단위는 청크로 나누어 처리해도 최종 결과를 합치는 데 문제가 없습니다. 이러한 작업의 특성을 파악하는 것이 최적화의 핵심입니다.
3. 튜닝 가능한 파라미터 도입: max_text_batch와 같이 사용자가 환경에 맞게 조절할 수 있는 파라미터를 제공함으로써, 모델의 유연성과 적용 범위를 넓힐 수 있습니다.
4. 정확성 검증의 중요성: VRAM 최적화 과정에서 모델의 정확성이 저하되지 않도록 철저한 검증(FP32/FP16 비교, IoU 측정 등)이 필수적입니다.

결론

Ultralytics SAM3 모델의 텍스트 프롬프팅 기능에 도입된 청크 기반 처리 방식은 VRAM 사용량을 크게 줄여 모델의 접근성을 높이는 중요한 개선입니다. 특히 대규모 프롬프트 세트를 다루는 애플리케이션에서 이 최적화는 큰 가치를 발휘할 것입니다. 성능 저하 없이 메모리 효율성을 높인 이 PR은 복잡한 딥러닝 모델을 실제 환경에 배포할 때 고려해야 할 중요한 최적화 전략을 잘 보여줍니다.

참고 자료

• https://github.com/ultralytics/ultralytics/blob/main/ultralytics/models/sam/sam3/sam3_image.py

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