Trellis 2がROCm 7.11上でAMD RX 9070 XTで正常に動作中

AMDハードウェアでTrellis 2を動作させる

開発者は、Linux Mint 22.3上でROCm 7.11を使用してAMD RX 9070 XT GPUでTrellis 2の実行に成功しました。これにより、AMDハードウェアでTrellis 2を実行しようとした際にユーザーが遭遇していたジオメトリのカットオフ、プレビューの失敗、その他のエラーといった一般的な問題が解決されます。

主要な問題と解決策

開発者は、ほとんどの失敗の原因となっていた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
));

結果とパフォーマンス

これらの修正により、開発者は画像から3Dへのパイプライン（法線なしのプレビュー描画と最終的なGLBエクスポートを含む）の動作に成功しました。21,204トークンのテスト画像では、開始からプレビュー生成までに約280秒かかりました。この実行では、すべてのサンプラーを20ステップに設定した1024解像度が使用されました。