Trellis 2가 AMD RX 9070 XT에서 ROCm 7.11로 성공적으로 실행 중입니다.

AMD 하드웨어에서 Trellis 2 실행하기

한 개발자가 Linux Mint 22.3에서 ROCm 7.11을 사용해 AMD RX 9070 XT GPU에서 Trellis 2를 성공적으로 실행했습니다. 이는 AMD 하드웨어에서 Trellis 2를 실행하려 할 때 사용자들이 겪는 지오메트리 절단, 미리보기 실패 및 기타 오류 문제를 해결합니다.

주요 문제점과 해결책

개발자는 대부분의 실패를 일으키는 두 가지 주요 문제를 확인했습니다:

1. 높은 N 텐서에서의 ROCm 불안정성

ROCm 연산은 큰 텐서에서 불안정해져 오버플로우나 NaN 값을 발생시킵니다. sparse 폴더의 linear.py에 있던 원본 코드는 다음과 같습니다:

def forward(self, input: VarLenTensor) -> VarLenTensor:
    return input.replace(super().forward(input.feats))

수정된 코드는 ROCm 문제를 피하기 위해 청크 처리 방식을 구현합니다:

ROCM_SAFE_CHUNK = 524_288
def rocm_safe_linear(feats: torch.Tensor, weight: torch.Tensor, bias=None) -> torch.Tensor:
    """F.linear with ROCm large-N chunking workaround."""
    N = feats.shape[0]
    if N <= ROCM_SAFE_CHUNK:
        return F.linear(feats, weight, bias)
    out = torch.empty(N, weight.shape[0], device=feats.device, dtype=feats.dtype)
    for s in range(0, N, ROCM_SAFE_CHUNK):
        e = min(s + ROCM_SAFE_CHUNK, N)
        out[s:e] = F.linear(feats[s:e], weight, bias)
    return out

def forward(self, input):
    feats = input.feats if hasattr(input, 'feats') else input
    out = rocm_safe_linear(feats, self.weight, self.bias)
    if hasattr(input, 'replace'):
        return input.replace(out)
    return out

2. CuMesh의 손상된 hipMemcpy2D

CuMesh의 hipMemcpy2D 함수는 정점과 면이 누락되거나 손상되는 문제를 일으켰습니다. 원본 CuMesh 초기화 코드는 다음과 같습니다:

void CuMesh::init(const torch::Tensor& vertices, const torch::Tensor& faces) {
    size_t num_vertices = vertices.size(0);
    size_t num_faces = faces.size(0);
    this->vertices.resize(num_vertices);
    this->faces.resize(num_faces);
    CUDA_CHECK(cudaMemcpy2D(
        this->vertices.ptr,
        sizeof(float3),
        vertices.data_ptr(),
        sizeof(float) * 3,
        sizeof(float) * 3,
        num_vertices,
        cudaMemcpyDeviceToDevice
    ));
    ...
}

수정된 코드는 2D 복사를 1D 버전으로 대체합니다:

CUDA_CHECK(cudaMemcpy(
    this->vertices.ptr,
    vertices.data_ptr(),
    num_vertices * sizeof(float3),
    cudaMemcpyDeviceToDevice
));

결과와 성능

이러한 수정으로 개발자는 미리보기 렌더링(노멀 없이)과 최종 GLB 내보내기를 포함한 이미지-3D 파이프라인을 성공적으로 작동시켰습니다. 21,204개의 토큰을 가진 테스트 이미지에서, 시작부터 미리보기 생성까지 약 280초가 소요되었습니다. 실행은 모든 샘플러를 20단계로 설정한 1024 해상도에서 이루어졌습니다.