# How to Fine-Tune Pretrained TRELLIS.2 Checkpoints at Higher Resolutions Using the `finetune_ckpt` Config

> Learn to fine-tune pretrained TRELLIS.2 checkpoints at higher resolutions. Update config settings, point to your weights, and run train.py for improved model performance.

- Repository: [Microsoft/TRELLIS.2](https://github.com/microsoft/TRELLIS.2)
- Tags: how-to-guide
- Published: 2026-08-04

---

**To fine-tune pretrained TRELLIS.2 checkpoints at higher resolutions, copy an existing configuration file, update the resolution parameters in both the decoder and dataset sections, set the `finetune_ckpt` field to point to your pretrained weights, and launch [`train.py`](https://github.com/microsoft/TRELLIS.2/blob/main/train.py) with the new config.**

The `microsoft/TRELLIS.2` repository provides pretrained 3-D VAE and flow checkpoints trained at resolutions like 512³. You can push these models to higher voxel resolutions—such as 1024³ or 1536³—without modifying the training code by leveraging the built-in `finetune_ckpt` configuration system.

## How the `finetune_ckpt` Configuration Works

The fine-tuning mechanism is implemented in [`trellis2/trainers/basic.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/trainers/basic.py). When you specify a `finetune_ckpt` dictionary in your JSON configuration, the `BasicTrainer` automatically loads the pretrained weights before the training loop begins.

### The Loading Logic in `BasicTrainer`

Inside `BasicTrainer.__init__` (lines 148–151 of [`basic.py`](https://github.com/microsoft/TRELLIS.2/blob/main/basic.py)), the trainer stores the `finetune_ckpt` argument. If no checkpoint directory is supplied via `load_dir`, the trainer immediately invokes `self.finetune_from(finetune_ckpt)`.

The `finetune_from()` method iterates over the dictionary of checkpoint paths, loads each `.pth` file using `torch.load(..., map_location=self.device)`, and copies the state dictionary into the corresponding model via `model.load_state_dict()`. After loading, the trainer proceeds with the standard training loop, but the model parameters are initialized from your pretrained weights rather than random initialization.

## Step-by-Step Resolution Upscaling

Follow this workflow to upscale a 512³ model to 1024³ or beyond.

### 1. Select a Base Configuration

Start with an existing configuration that contains the architecture definitions you need. The repository includes reference configs such as [`configs/scvae/shape_vae_next_dc_f16c32_fp16_ft_512.json`](https://github.com/microsoft/TRELLIS.2/blob/main/configs/scvae/shape_vae_next_dc_f16c32_fp16_ft_512.json), which defines the `FlexiDualGridVaeEncoder` and `FlexiDualGridVaeDecoder` used for shape auto-encoding.

### 2. Adjust Resolution Parameters

Copy the base JSON to a new file (e.g., [`shape_vae_next_dc_f16c32_fp16_ft_1024.json`](https://github.com/microsoft/TRELLIS.2/blob/main/shape_vae_next_dc_f16c32_fp16_ft_1024.json)) and update the following fields:

- **`models.decoder.args.resolution`**: Change from `512` to `1024` (or your target resolution).
- **`dataset.args.resolution`**: Match the decoder resolution to ensure the data loader provides correctly sized voxel grids.
- **`models.encoder.args.resolution`**: Typically handled implicitly by the decoder configuration, but verify your specific model class requirements.

### 3. Point to Pretrained Weights

Add a `finetune_ckpt` dictionary inside the `trainer.args` section. Map each model component to its corresponding checkpoint file from your previous training run:

```json
"finetune_ckpt": {
  "encoder": "output/shape_vae_512/ckpts/shape_encoder_512.pth",
  "decoder": "output/shape_vae_512/ckpts/shape_decoder_512.pth"
}

```

### 4. Launch Training

Run [`train.py`](https://github.com/microsoft/TRELLIS.2/blob/main/train.py) with your new configuration file. The trainer will automatically detect the `finetune_ckpt` entries, load the weights, and begin training at the higher resolution.

## Key Configuration Parameters for Higher Resolution

When upscaling, adjust these hyperparameters to ensure stable fine-tuning:

| Parameter | 512³ (Reference) | 1024³ (Example) |
|-----------|------------------|-----------------|
| `models.decoder.args.resolution` | `512` | `1024` |
| `dataset.args.resolution` | `512` | `1024` |
| `trainer.args.max_steps` | `1_000_000` | `200_000` (shorter fine-tuning run) |
| `trainer.args.lr` | `1e-5` | `5e-6` (lower learning rate to preserve features) |

Because the architecture supports flexible resolutions, the loaded weights reshape automatically where possible. You do not need to alter the model code in [`trellis2/modules/vae.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/vae.py).

## Code Example: 1024³ Fine-Tuning Config

Below is a complete configuration file adapted from the 512³ base. Note the updated resolution fields and the `finetune_ckpt` dictionary pointing to the 512³ weights:

```json
{
  "models": {
    "encoder": {
      "name": "FlexiDualGridVaeEncoder",
      "args": {
        "model_channels": [64, 128, 256, 512, 1024],
        "latent_channels": 32,
        "num_blocks": [0, 4, 8, 16, 4],
        "block_type": ["SparseConvNeXtBlock3d", "SparseConvNeXtBlock3d", "SparseConvNeXtBlock3d", "SparseConvNeXtBlock3d", "SparseConvNeXtBlock3d"],
        "down_block_type": ["SparseResBlockS2C3d", "SparseResBlockS2C3d", "SparseResBlockS2C3d", "SparseResBlockS2C3d"],
        "block_args": [{ "use_checkpoint": true }, { "use_checkpoint": true }, { "use_checkpoint": true }, { "use_checkpoint": true }, { "use_checkpoint": true }],
        "use_fp16": true
      }
    },
    "decoder": {
      "name": "FlexiDualGridVaeDecoder",
      "args": {
        "resolution": 1024,
        "model_channels": [1024, 512, 256, 128, 64],
        "latent_channels": 32,
        "num_blocks": [4, 16, 8, 4, 0],
        "block_type": ["SparseConvNeXtBlock3d", "SparseConvNeXtBlock3d", "SparseConvNeXtBlock3d", "SparseConvNeXtBlock3d", "SparseConvNeXtBlock3d"],
        "up_block_type": ["SparseResBlockC2S3d", "SparseResBlockC2S3d", "SparseResBlockC2S3d", "SparseResBlockC2S3d"],
        "block_args": [{ "use_checkpoint": true }, { "use_checkpoint": true }, { "use_checkpoint": true }, { "use_checkpoint": true }, { "use_checkpoint": true }],
        "use_fp16": true
      }
    }
  },
  "dataset": {
    "name": "FlexiDualGridDataset",
    "args": {
      "resolution": 1024,
      "max_active_voxels": 1000000,
      "max_num_faces": 1000000,
      "min_aesthetic_score": 4.5
    }
  },
  "trainer": {
    "name": "ShapeVaeTrainer",
    "args": {
      "max_steps": 200000,
      "batch_size_per_gpu": 2,
      "optimizer": { "name": "AdamW", "args": { "lr": 5e-6, "weight_decay": 0.0 } },
      "finetune_ckpt": {
        "encoder": "output/shape_vae_512/ckpts/shape_encoder_512.pth",
        "decoder": "output/shape_vae_512/ckpts/shape_decoder_512.pth"
      },
      "fp16_mode": "inflat_all",
      "grad_clip": { "name": "AdaptiveGradClipper", "args": { "max_norm": 1.0, "clip_percentile": 95 } }
    }
  }
}

```

## Running the Training

Save your configuration to [`configs/scvae/shape_vae_next_dc_f16c32_fp16_ft_1024.json`](https://github.com/microsoft/TRELLIS.2/blob/main/configs/scvae/shape_vae_next_dc_f16c32_fp16_ft_1024.json), then execute:

```bash
python train.py \
  --config configs/scvae/shape_vae_next_dc_f16c32_fp16_ft_1024.json \
  --output_dir results/shape_vae_1024 \
  --data_dir '{"ObjaverseXL_sketchfab": {"base": "datasets/ObjaverseXL_sketchfab", "dual_grid": "datasets/ObjaverseXL_sketchfab/dual_grid_1024"}}'

```

This command instructs [`train.py`](https://github.com/microsoft/TRELLIS.2/blob/main/train.py) to initialize the `ShapeVaeTrainer`, load the encoder and decoder weights from the paths specified in `finetune_ckpt`, and continue training on 1024³ dual-grid data.

## Summary

- **No code changes required**: The `finetune_ckpt` mechanism in [`trellis2/trainers/basic.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/trainers/basic.py) handles weight loading automatically when `load_dir` is `None`.
- **Edit configuration only**: Update `models.decoder.args.resolution`, `dataset.args.resolution`, and point `finetune_ckpt` to your existing `.pth` files.
- **Adjust training duration**: Reduce `max_steps` and learning rate (`lr`) when fine-tuning to avoid overwriting learned features.
- **Scalable workflow**: Repeat this process to reach 1536³ or higher by incrementing the resolution fields and updating the data path.

## Frequently Asked Questions

### Can I fine-tune from a checkpoint while changing both the resolution and the latent dimension?

No, you should keep the `latent_channels` and model channel dimensions identical to the pretrained checkpoint. The `finetune_from()` method performs strict state-dict loading via `model.load_state_dict()`, so architectural mismatches will raise errors. Only resolution-related parameters (which affect the spatial grid size but not the channel counts) can be safely modified.

### Why should I reduce the learning rate when fine-tuning at higher resolutions?

Lower resolutions (512³) capture coarse geometric features, while higher resolutions (1024³) require refinement of fine details. A smaller learning rate—such as `5e-6` instead of `1e-5`—preserves the coarse features learned at the lower resolution while allowing gradual adaptation to the finer voxel grid, preventing catastrophic forgetting.

### Does the `finetune_ckpt` option load optimizer states or only model weights?

The `finetune_ckpt` mechanism loads only model weights. According to the implementation in [`trellis2/trainers/basic.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/trainers/basic.py), the `finetune_from()` method uses `torch.load()` and `load_state_dict()` on the models, but does not restore optimizer states, EMA buffers, or step counters. This allows you to restart training with a fresh optimizer configuration suitable for the new resolution.

### Can I use this method to fine-tune the flow matching models, or only the VAE?

The `finetune_ckpt` configuration works for any trainer inheriting from `BasicTrainer`, including flow matching models. Ensure your config defines the appropriate model classes (e.g., flow-based architectures) and that the checkpoint paths in `finetune_ckpt` match the model keys expected by your specific trainer implementation.