fix: construct BF16 checkpoints in FP16 on pre-Ampere GPUs

README.md

@@ -76,7 +76,7 @@ Data processing is streamlined for instant conversions that are fully **renderin
     - By adding `--new-env`, a new conda environment named `trellis2` will be created. If you want to use an existing conda environment, please remove this flag.
     - By default the `trellis2` environment will use pytorch 2.6.0 with CUDA 12.4. If you want to use a different version of CUDA, you can remove the `--new-env` flag and manually install the required dependencies. Refer to [PyTorch](https://pytorch.org/get-started/previous-versions/) for the installation command.
     - If you have multiple CUDA Toolkit versions installed, `CUDA_HOME` should be set to the correct version before running the command. For example, if you have CUDA Toolkit 12.4 and 13.0 installed, you can run `export CUDA_HOME=/usr/local/cuda-12.4` before running the command.
-    - By default, the code uses the `flash-attn` backend for attention. For GPUs do not support `flash-attn` (e.g., NVIDIA V100), you can install `xformers` manually and set the `ATTN_BACKEND` environment variable to `xformers` before running the code. See the [Minimal Example](#minimal-example) for more details.
+    - By default, the code uses the `flash-attn` backend for attention. For GPUs do not support `flash-attn` (e.g., NVIDIA V100), you can install `xformers` manually and set the `ATTN_BACKEND` environment variable to `xformers` before running the code. See the [Minimal Example](#minimal-example) for more details. When the current CUDA device has compute capability below 8.0 (pre-Ampere), BF16 checkpoints are automatically constructed in FP16. On mixed-GPU systems, select the intended CUDA device before loading (e.g., with `torch.cuda.set_device`). FP16 may overflow or produce numerically unstable results, and older GPUs may still encounter separate backend failures, including Triton/FlexGEMM or spconv errors.
     - The installation may take a while due to the large number of dependencies. Please be patient. If you encounter any issues, you can try to install the dependencies one by one, specifying one flag at a time.
     - If you encounter any issues during the installation, feel free to open an issue or contact us.
 

trellis2/models/__init__.py

@@ -46,6 +46,8 @@ def from_pretrained(path: str, **kwargs):
     """
     import os
     import json
+    import warnings
+    import torch
     from safetensors.torch import load_file
     is_local = os.path.exists(f"{path}.json") and os.path.exists(f"{path}.safetensors")
 
@@ -62,6 +64,14 @@ def from_pretrained(path: str, **kwargs):
 
     with open(config_file, 'r') as f:
         config = json.load(f)
+    if config['args'].get('dtype') in ('bf16', 'bfloat16') and torch.cuda.is_available():
+        if torch.cuda.get_device_capability()[0] < 8:
+            config['args']['dtype'] = 'float16'
+            warnings.warn(
+                "The current CUDA device is pre-Ampere; loading this BF16 checkpoint in float16. "
+                "FP16 may overflow and produce numerically unstable results.",
+                stacklevel=2,
+            )
     model = __getattr__(config['name'])(**config['args'], **kwargs)
     model.load_state_dict(load_file(model_file), strict=False)