Trellis 2 fonctionne avec succès sur ROCm 7.11 avec une AMD RX 9070 XT

Faire fonctionner Trellis 2 sur du matériel AMD

Un développeur a réussi à exécuter Trellis 2 sur une carte graphique AMD RX 9070 XT en utilisant ROCm 7.11 sous Linux Mint 22.3. Cela résout les problèmes courants où les utilisateurs rencontraient des coupures géométriques, des échecs de prévisualisation et d'autres erreurs en essayant de faire fonctionner Trellis 2 sur du matériel AMD.

Problèmes clés et solutions

Le développeur a identifié deux problèmes principaux qui causaient la plupart des échecs :

1. Instabilité de ROCm avec les tenseurs de grande dimension N

Les opérations ROCm deviennent instables avec les grands tenseurs, provoquant des dépassements de capacité ou des valeurs NaN. Le code original dans linear.py du dossier sparse utilisait :

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

La solution implémente un traitement par morceaux pour éviter les problèmes ROCm :

ROCM_SAFE_CHUNK = 524_288
def rocm_safe_linear(feats: torch.Tensor, weight: torch.Tensor, bias=None) -> torch.Tensor:
    """F.linear avec un contournement par morceaux pour ROCm avec de grandes valeurs de N."""
    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. Fonction hipMemcpy2D défectueuse dans CuMesh

La fonction hipMemcpy2D dans CuMesh provoquait la disparition ou la corruption des sommets et des faces. L'initialisation originale de CuMesh utilisait :

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

La solution remplace la copie 2D par une version 1D :

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

Résultats et performances

Avec ces correctifs, le développeur a réussi à faire fonctionner le pipeline image-3D, y compris le rendu de prévisualisation (sans normales) et l'export final en GLB. Sur une image de test avec 21 204 tokens, le processus a pris environ 280 secondes du début à la génération de la prévisualisation. L'exécution utilisait une résolution de 1024 avec tous les échantillonneurs réglés sur 20 étapes.