From 0e1efdbef3f6c934f9551bbcf9832b9db6930cd4 Mon Sep 17 00:00:00 2001 From: Matt Van Horn <455140+mvanhorn@users.noreply.github.com> Date: Sat, 10 Oct 2026 22:03:57 -0700 Subject: [PATCH] fix: construct BF16 checkpoints in FP16 on pre-Ampere GPUs For checkpoints configured as bf16 or bfloat16, from_pretrained now substitutes float16 before model construction when CUDA is available and the current device has compute capability below 8.0. The fallback emits a warning about FP16 overflow and numerical instability. The README explains the fallback, its dependence on the current CUDA device on mixed-GPU systems, and the separate backend failures older GPUs may still encounter. Fixes #23 --- README.md | 2 +- trellis2/models/__init__.py | 10 ++++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 1c0ad06b..adc3b9f4 100644 --- a/README.md +++ b/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. diff --git a/trellis2/models/__init__.py b/trellis2/models/__init__.py index d4fed035..65808636 100644 --- a/trellis2/models/__init__.py +++ b/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)