Trellis 2 läuft erfolgreich auf ROCm 7.11 mit AMD RX 9070 XT

Trellis 2 auf AMD-Hardware zum Laufen bringen

Ein Entwickler hat Trellis 2 erfolgreich auf einer AMD RX 9070 XT GPU mit ROCm 7.11 unter Linux Mint 22.3 ausgeführt. Dies behebt häufige Probleme, bei denen Benutzer Geometrieabbrüche, Vorschaufehler und andere Fehler erlebten, wenn sie versuchten, Trellis 2 auf AMD-Hardware auszuführen.

Hauptprobleme und Lösungen

Der Entwickler identifizierte zwei Hauptprobleme, die die meisten Fehler verursachten:

1. ROCm-Instabilität bei hohen N-Tensoren

ROCm-Operationen werden bei großen Tensoren instabil und verursachen Überläufe oder NaN-Werte. Der ursprüngliche Code in linear.py im sparse-Ordner verwendete:

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

Die Lösung implementiert eine blockweise Verarbeitung, um ROCm-Probleme zu vermeiden:

ROCM_SAFE_CHUNK = 524_288
def rocm_safe_linear(feats: torch.Tensor, weight: torch.Tensor, bias=None) -> torch.Tensor:
    """F.linear mit ROCm-Chunking-Workaround für große 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. Fehlerhafte hipMemcpy2D in CuMesh

Die Funktion hipMemcpy2D in CuMesh verursachte, dass Vertices und Faces verloren gingen oder beschädigt wurden. Die ursprüngliche CuMesh-Initialisierung verwendete:

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

Die Lösung ersetzt den 2D-Kopierbefehl durch eine 1D-Version:

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

Ergebnisse und Leistung

Mit diesen Korrekturen konnte der Entwickler die Bild-zu-3D-Pipeline erfolgreich zum Laufen bringen, einschließlich Vorschau-Rendering (ohne Normalen) und finalem GLB-Export. Bei einem Testbild mit 21.204 Tokens dauerte der Prozess etwa 280 Sekunden vom Start bis zur Vorschau-Erstellung. Der Lauf verwendete eine Auflösung von 1024 mit allen Samplern auf 20 Schritte eingestellt.