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

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 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. 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), 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, 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) 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:

"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 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.

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:

{
  "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, then execute:

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 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 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, 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.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}