Trellis 2 Executando com Sucesso no ROCm 7.11 com AMD RX 9070 XT

Fazendo o Trellis 2 Funcionar em Hardware AMD

Um desenvolvedor executou com sucesso o Trellis 2 em uma GPU AMD RX 9070 XT usando ROCm 7.11 no Linux Mint 22.3. Isso resolve problemas comuns em que os usuários encontravam corte de geometria, falhas na visualização prévia e outros erros ao tentar executar o Trellis 2 em hardware AMD.

Problemas Principais e Soluções

O desenvolvedor identificou dois problemas principais que causavam a maioria das falhas:

1. Instabilidade do ROCm com Tensores N Altos

As operações do ROCm ficam instáveis com tensores grandes, causando estouros ou valores NaN. O código original em linear.py na pasta sparse usava:

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

A correção implementa processamento em blocos para evitar problemas no ROCm:

ROCM_SAFE_CHUNK = 524_288
def rocm_safe_linear(feats: torch.Tensor, weight: torch.Tensor, bias=None) -> torch.Tensor:
    """F.linear com solução alternativa de divisão em blocos para ROCm com N grande."""
    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. Função hipMemcpy2D Quebrada no CuMesh

A função hipMemcpy2D no CuMesh estava fazendo vértices e faces desaparecerem ou ficarem corrompidos. A inicialização original do CuMesh usava:

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
    ));
    ...
}

A correção substitui a cópia 2D por uma versão 1D:

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

Resultados e Desempenho

Com essas correções, o desenvolvedor conseguiu fazer o pipeline de imagem para 3D funcionar, incluindo renderização de visualização prévia (sem normais) e exportação final GLB. Em uma imagem de teste com 21.204 tokens, o processo levou aproximadamente 280 segundos do início até a geração da visualização prévia. A execução usou resolução 1024 com todos os amostradores configurados para 20 passos.