Skip to content

offloading VAE to CPU makes inference latency unacceptable #99

Description

@CCSCovenant

Problem

The README currently documents single-GPU inference with offload mode enabled, specifically saying that VAE and text_encoder are offloaded to CPU for the RoboTwin and i2av paths:

  • RoboTwin: approximately 24GB VRAM with offload mode enabled, VAE and text_encoder offloaded to CPU.
  • i2av: approximately 18GB VRAM with offload mode enabled, VAE and text_encoder offloaded to CPU.

In our tests, following this broad offload behavior makes inference latency unacceptable for deployment/control use. VAE CPU offload is the critical blocker: the model initializes, but a single inference chunk does not complete within a 240s timeout.

Environment

  • GPU: RTX 4090 24GB
  • PyTorch: 2.9.0+cu128
  • CUDA runtime: 12.8
  • Attention backend: flashattn
  • Input size: 256x256
  • frame_chunk_size: 4
  • num_inference_steps: 5
  • action_num_inference_steps: 10

Observed results

Case 1: VAE + text_encoder offloaded to CPU

Configuration:

enable_offload=true
enable_vae_offload=true
enable_text_encoder_offload=true

Timing:

model_init: 2.462s
reset: 86.625s
one_chunk_inference: started but did not finish within 240s

This is not usable for real-time or closed-loop robot control.

Case 2: keep VAE on GPU, offload only text_encoder

Configuration:

enable_offload=true
enable_vae_offload=false
enable_text_encoder_offload=true

Timing:

model_load: 2.421s
reset: 89.740s
one_chunk_inference: 1.623s

Memory peak during one chunk:

max_allocated: 20210.8 MiB
max_reserved: 20292.0 MiB

A 32-chunk wrapper run with VAE on GPU completed 32/32 requests successfully. After the first reset-heavy request, subsequent chunks were about 1.4-2.0s each.

Expected behavior / request

Please fix readme instruction so that For 24GB GPUs, keeping the VAE resident on GPU appears to be necessary for practical inference.

Related: #41 also points out that the original code offloads VAE to CPU.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions