Add low-memory relay mode for T2I inference

.gitignore

@@ -32,6 +32,7 @@ results/
 downloads/
 temps/
 datasets/
+pr-evidence/
 
 #
 outputs/

README.md

@@ -127,7 +127,7 @@ We are actively updating and improving this repository. If you find any bugs or
 ### Recommended Environment
 
 - **Software:** Python 3.10+, CUDA 12.4+ (required)
-- **Hardware:** A GPU with at least 40GB VRAM is required for inference
+- **Hardware:** Default parallel inference requires a GPU with at least 40GB VRAM. For text-to-image only, `MEMORY_MODE=relay` lowers peak GPU memory by loading the UND tower, GEN tower, and VAE only for the phase that needs them.
 
 We have tested the following dependency combinations on NVIDIA A100:
 
@@ -240,6 +240,24 @@ bash inference_lance.sh \
   --SAVE_PATH_GEN results/t2i
 ```
 
+##### Low-Memory Text-to-Image Relay
+
+`MEMORY_MODE=parallel` is the default and fastest path. For text-to-image inference with KV-cache enabled, `MEMORY_MODE=relay` runs the same T2I pipeline while relaying the UND tower, GEN tower, and VAE through GPU memory one phase at a time. Relay is slower, but it can substantially reduce peak GPU memory and is intended to produce bit-identical images to `parallel` on the same hardware and software stack with the same seed and settings.
+
+```bash
+bash inference_lance.sh \
+  --TASK_NAME t2i \
+  --MODEL_PATH downloads/Lance_3B \
+  --RESOLUTION image_768res \
+  --VIDEO_HEIGHT 768 \
+  --VIDEO_WIDTH 768 \
+  --USE_KVCACHE true \
+  --MEMORY_MODE relay \
+  --SAVE_PATH_GEN results/t2i_relay
+```
+
+Current relay support is limited to `t2i` with `--USE_KVCACHE true`. Use `--RELAY_MEMORY_LOG true` to print CUDA allocator summaries at relay phase boundaries. To compare `parallel` and `relay` outputs on your machine, run `python tools/compare_memory_modes_t2i.py --model-path downloads/Lance_3B`.
+
 ##### Video Editing
 
 ```bash
@@ -323,6 +341,8 @@ You can configure the following hyperparameters at the top of the `inference_lan
 | `RESOLUTION` | `"video_480p"` | Base resolution preset (`image_768res` or `video_480p`). |
 | `CONFIG_PATH` | `""` | Optional path to a custom validation JSON/JSONL file. When empty, the task default example config is used. |
 | `ENHANCE_PROMPT` | `false` | Optional T2V/I2V prompt rewrite switch. T2V uses text-only rewrite; I2V uses text plus the input image. Prompt enhancement generally improves generation quality. This option requires `openai==2.26.0`; it is included in `requirements.txt`, or install it manually with `pip install openai==2.26.0`. Configure `API_KEY`, `MODEL_NAME`, and `BASE_URL` in `common/utils/caption_rewrite.py` before setting this to `true`; without a valid rewrite config, we recommend writing prompts in the style of the provided examples. |
+| `MEMORY_MODE` | `parallel` | GPU memory policy. Use `parallel` for the default fastest path, or `relay` for low-memory T2I with KV-cache. |
+| `RELAY_MEMORY_LOG` | `false` | Print CUDA allocator summaries at relay phase boundaries when using `MEMORY_MODE=relay`. |
 
 </details>
 

config/config_factory.py

@@ -306,6 +306,8 @@ class InferenceArguments(TrainingArguments):
     system_prompt_type:         str = "SP0"  # options: SP1, SP2 ...
     use_KVcache:                bool = False
     enhance_prompt:             bool = False  # Rewrite T2V prompts before inference when enabled.
+    memory_mode:                str = "parallel"  # parallel | relay. Relay is t2i + KV-cache only.
+    relay_memory_log:           bool = False  # Print CUDA allocator usage at relay phase boundaries.
 
 
 @dataclass

inference_lance.sh

@@ -15,6 +15,8 @@ VALIDATION_DATA_SEED=${VALIDATION_DATA_SEED:-42}
 CFG_TEXT_SCALE=${CFG_TEXT_SCALE:-4.0}
 USE_KVCACHE=${USE_KVCACHE:-true}
 ENHANCE_PROMPT=${ENHANCE_PROMPT:-false}
+MEMORY_MODE=${MEMORY_MODE:-parallel}       # parallel | relay
+RELAY_MEMORY_LOG=${RELAY_MEMORY_LOG:-false}
 
 NUM_FRAMES=${NUM_FRAMES:-50}             # max: 121 frames, unused for image tasks
 VIDEO_HEIGHT=${VIDEO_HEIGHT:-768}        # unused for editing
@@ -39,6 +41,8 @@ while [[ $# -gt 0 ]]; do
         --CFG_TEXT_SCALE) CFG_TEXT_SCALE="$2"; shift 2 ;;
         --USE_KVCACHE) USE_KVCACHE="$2"; shift 2 ;;
         --ENHANCE_PROMPT) ENHANCE_PROMPT="$2"; shift 2 ;;
+        --MEMORY_MODE) MEMORY_MODE="$2"; shift 2 ;;
+        --RELAY_MEMORY_LOG) RELAY_MEMORY_LOG="$2"; shift 2 ;;
 
         --NUM_FRAMES) NUM_FRAMES="$2"; shift 2 ;;
         --VIDEO_HEIGHT) VIDEO_HEIGHT="$2"; shift 2 ;;
@@ -48,13 +52,14 @@ while [[ $# -gt 0 ]]; do
         --SAVE_PATH_GEN) SAVE_PATH_GEN="$2"; shift 2 ;;
 
         -h|--help)
-            echo "Usage: bash inference_lance_my.sh [OPTIONS]"
+            echo "Usage: bash inference_lance.sh [OPTIONS]"
             echo ""
             echo "Example:"
-            echo "  bash inference_lance_my.sh --TASK_NAME t2i --MODEL_PATH downloads/Lance_3B --RESOLUTION image_768res"
-            echo "  bash inference_lance_my.sh --TASK_NAME image_edit --CONFIG_PATH config.json"
-            echo "  bash inference_lance_my.sh --TASK_NAME t2v --ENHANCE_PROMPT true"
-            echo "  bash inference_lance_my.sh --TASK_NAME i2v --ENHANCE_PROMPT true"
+            echo "  bash inference_lance.sh --TASK_NAME t2i --MODEL_PATH downloads/Lance_3B --RESOLUTION image_768res"
+            echo "  bash inference_lance.sh --TASK_NAME image_edit --CONFIG_PATH config.json"
+            echo "  bash inference_lance.sh --TASK_NAME t2v --ENHANCE_PROMPT true"
+            echo "  bash inference_lance.sh --TASK_NAME i2v --ENHANCE_PROMPT true"
+            echo "  bash inference_lance.sh --TASK_NAME t2i --MODEL_PATH downloads/Lance_3B --MEMORY_MODE relay --RELAY_MEMORY_LOG true"
             exit 0
             ;;
 
@@ -111,6 +116,8 @@ echo "  - cfg_text_scale: ${CFG_TEXT_SCALE}"
 echo "  - num_frames: ${NUM_FRAMES}"
 echo "  - use_KVcache: ${USE_KVCACHE}"
 echo "  - enhance_prompt: ${ENHANCE_PROMPT}"
+echo "  - memory_mode: ${MEMORY_MODE}"
+echo "  - relay_memory_log: ${RELAY_MEMORY_LOG}"
 echo "================================================"
 echo ""
 
@@ -157,6 +164,8 @@ accelerate launch \
     --cfg_text_scale        $CFG_TEXT_SCALE \
     --use_KVcache           "$USE_KVCACHE" \
     --enhance_prompt        "$ENHANCE_PROMPT" \
+    --memory_mode           "$MEMORY_MODE" \
+    --relay_memory_log      "$RELAY_MEMORY_LOG" \
     "${CONFIG_ARGS[@]}"
 
 echo ""

modeling/lance/lance.py

@@ -131,11 +131,141 @@ def __init__(
 
         self.config = config
         self.training_args: TrainingArguments = kwargs.get("training_args")
+        self._relay_memory_enabled = False
+        self._relay_offload_device = torch.device("cpu")
+        self._relay_log_memory = False
 
     def update_tokenizer(self, tokenizer):
         self.tokenizer: Qwen2Tokenizer = tokenizer
         self.vocab_size_efficient = len(tokenizer)
 
+    def configure_relay_memory(self, enabled: bool, offload_device: str = "cpu", log_memory: bool = False):
+        self._relay_memory_enabled = enabled
+        self._relay_offload_device = torch.device(offload_device)
+        self._relay_log_memory = log_memory
+        self._set_relay_stream_und(False)
+
+    def _set_relay_stream_und(self, enabled: bool):
+        qwen_model = self.language_model.model
+        qwen_model._relay_stream_und = enabled
+        qwen_model._relay_offload_device = self._relay_offload_device
+        qwen_model._relay_log_memory = self._relay_log_memory
+
+    def _relay_cuda_log(self, stage_name: str, device=None):
+        if not self._relay_log_memory or not torch.cuda.is_available():
+            return
+        if device is None:
+            device = torch.cuda.current_device()
+        torch.cuda.synchronize(device)
+        allocated = torch.cuda.memory_allocated(device) / (1024 ** 3)
+        reserved = torch.cuda.memory_reserved(device) / (1024 ** 3)
+        peak = torch.cuda.max_memory_allocated(device) / (1024 ** 3)
+        self.log_rank0(f"[relay] {stage_name}: allocated={allocated:.2f}GiB reserved={reserved:.2f}GiB peak={peak:.2f}GiB")
+
+    def _relay_move_modules(self, modules, device, dtype=None):
+        seen = set()
+        for module in modules:
+            if module is None or not isinstance(module, nn.Module):
+                continue
+            module_id = id(module)
+            if module_id in seen:
+                continue
+            seen.add(module_id)
+            if dtype is None:
+                module.to(device=device)
+            else:
+                module.to(device=device, dtype=dtype)
+
+    @staticmethod
+    def _relay_named_child(module: nn.Module, name: str):
+        return getattr(module, name, None)
+
+    def _relay_und_tower_modules(self):
+        qwen_model = self.language_model.model
+        for layer in qwen_model.layers:
+            attn = layer.self_attn
+            for attr in ("q_proj", "k_proj", "v_proj", "o_proj", "q_norm", "k_norm"):
+                yield self._relay_named_child(attn, attr)
+            for attr in ("input_layernorm", "post_attention_layernorm", "mlp"):
+                yield self._relay_named_child(layer, attr)
+        yield self._relay_named_child(qwen_model, "norm")
+
+    def _relay_gen_tower_modules(self):
+        qwen_model = self.language_model.model
+        for layer in qwen_model.layers:
+            attn = layer.self_attn
+            for attr in ("q_proj_moe_gen", "k_proj_moe_gen", "v_proj_moe_gen", "o_proj_moe_gen", "q_norm_moe_gen", "k_norm_moe_gen"):
+                yield self._relay_named_child(attn, attr)
+            for attr in ("input_layernorm_moe_gen", "post_attention_layernorm_moe_gen", "mlp_moe_gen"):
+                yield self._relay_named_child(layer, attr)
+        yield self._relay_named_child(qwen_model, "norm_moe_gen")
+
+    def _relay_shared_prefill_modules(self):
+        qwen_model = self.language_model.model
+        yield self._relay_named_child(qwen_model, "embed_tokens")
+        yield self._relay_named_child(qwen_model, "rotary_emb")
+        yield self._relay_named_child(self, "time_embedder")
+        yield self._relay_named_child(self, "vae2llm")
+        yield self._relay_named_child(self, "llm2vae")
+        yield self._relay_named_child(self, "latent_pos_embed")
+
+    def _relay_shared_gen_modules(self):
+        qwen_model = self.language_model.model
+        yield self._relay_named_child(qwen_model, "rotary_emb")
+        yield self._relay_named_child(self, "time_embedder")
+        yield self._relay_named_child(self, "vae2llm")
+        yield self._relay_named_child(self, "llm2vae")
+        yield self._relay_named_child(self, "latent_pos_embed")
+
+    def relay_prepare_prefill(self, device, dtype):
+        if not self._relay_memory_enabled:
+            return
+        self._set_relay_stream_und(False)
+        self._relay_move_modules(self._relay_gen_tower_modules(), self._relay_offload_device)
+        self._relay_move_modules(self._relay_und_tower_modules(), device, dtype=dtype)
+        self._relay_move_modules(self._relay_shared_prefill_modules(), device, dtype=dtype)
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        self._relay_cuda_log("prefill modules ready", device)
+
+    def relay_switch_to_gen(self, device, dtype):
+        if not self._relay_memory_enabled:
+            return
+        self._relay_move_modules(self._relay_und_tower_modules(), self._relay_offload_device)
+        self._relay_move_modules([self.language_model.model.embed_tokens, self.vit_model if hasattr(self, "vit_model") else None, self.connector if hasattr(self, "connector") else None], self._relay_offload_device)
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        self._relay_cuda_log("UND tower offloaded", device)
+        self._relay_move_modules(self._relay_gen_tower_modules(), device, dtype=dtype)
+        self._relay_move_modules(self._relay_shared_gen_modules(), device, dtype=dtype)
+        self._set_relay_stream_und(True)
+        self._relay_cuda_log("GEN tower ready", device)
+
+    def relay_offload_all(self):
+        if not self._relay_memory_enabled:
+            self.to("cpu")
+            return
+        self._set_relay_stream_und(False)
+        self._relay_move_modules(self._relay_und_tower_modules(), self._relay_offload_device)
+        self._relay_move_modules(self._relay_gen_tower_modules(), self._relay_offload_device)
+        self._relay_move_modules(self._relay_shared_prefill_modules(), self._relay_offload_device)
+        self._relay_move_modules([self.language_model.lm_head], self._relay_offload_device)
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+        self._relay_cuda_log("all Lance modules offloaded")
+
+    def relay_discard_all(self):
+        if not self._relay_memory_enabled:
+            self.to_empty(device=torch.device("meta"))
+            return
+        self._set_relay_stream_und(False)
+        self.to_empty(device=torch.device("meta"))
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+        self._relay_cuda_log("all Lance modules discarded")
+
     def process_attention_mask(self, current_attn_modes, current_split_lens, current_seq_len, device, BLOCK_SIZE=128):
         current_attn_modes_ = ["full" if mode_ in ["full_noise", "full_noise_target"] else mode_ for mode_ in current_attn_modes]
         sparse_mask = create_sparse_mask(current_seq_len, current_split_lens, current_attn_modes_, device)
@@ -1393,6 +1523,7 @@ def validation_gen_KVcache(
         index: str = "",
         **kwargs,
     ):
+        relay_enabled = kwargs.get("memory_mode", "parallel") == "relay" and self._relay_memory_enabled
         cfg_vision_scale = cfg_vit_scale
         pt, ph, pw = self.latent_patch_size
         index_dtype = val_packed_text_ids.dtype
@@ -1433,6 +1564,9 @@ def validation_gen_KVcache(
             if gen_idx >= max_samples:
                 break
 
+            if relay_enabled:
+                self.relay_prepare_prefill(device, dtype)
+
             # 1. Get slice information for the current sample in the batch.
             sample_start_idx = cu_sample_lens[i_sample]
             sample_end_idx = cu_sample_lens[i_sample + 1]
@@ -1614,6 +1748,8 @@ def validation_gen_KVcache(
 
                 current_cond_start = current_cond_end
 
+            if relay_enabled:
+                self.relay_switch_to_gen(device, dtype)
 
             for _ in range(1):
                 timestep = torch.zeros(x_t.shape[0], device=x_t.device)

modeling/lance/qwen2_navit.py

@@ -592,6 +592,24 @@ def __init__(
         self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.post_attention_layernorm_moe_gen = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
+    def _relay_und_modules(self):
+        attn = self.self_attn
+        for attr in ("q_proj", "k_proj", "v_proj", "o_proj", "q_norm", "k_norm"):
+            module = getattr(attn, attr, None)
+            if isinstance(module, nn.Module):
+                yield module
+        for attr in ("input_layernorm", "post_attention_layernorm", "mlp"):
+            module = getattr(self, attr, None)
+            if isinstance(module, nn.Module):
+                yield module
+
+    def _relay_move_und_branch(self, device, dtype=None):
+        for module in self._relay_und_modules():
+            if dtype is None:
+                module.to(device=device)
+            else:
+                module.to(device=device, dtype=dtype)
+
     def forward(self, *args, **kwargs):
         if self.training or kwargs.get("mode_forward") == "validation":
             return self.forward_train(*args, **kwargs)
@@ -660,6 +678,17 @@ def forward_inference(
         **kwargs
     ) -> BaseNavitOutputWithPast:
 
+        relay_stream_und = bool(kwargs.get("relay_stream_und", False))
+        relay_offload_device = kwargs.get("relay_offload_device", torch.device("cpu"))
+        relay_stream_this_layer = (
+            relay_stream_und
+            and mode == "gen"
+            and packed_text_indexes is not None
+            and packed_text_indexes.numel() > 0
+        )
+        if relay_stream_this_layer:
+            self._relay_move_und_branch(packed_query_sequence.device, dtype=packed_query_sequence.dtype)
+
         residual = packed_query_sequence
         if mode == "und":
             packed_query_sequence = self.input_layernorm(packed_query_sequence)
@@ -704,6 +733,10 @@ def forward_inference(
             packed_query_sequence = packed_query_sequence_
 
         packed_query_sequence = residual + packed_query_sequence
+        if relay_stream_this_layer:
+            self._relay_move_und_branch(relay_offload_device)
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
         return packed_query_sequence, past_key_values
 
 
@@ -919,6 +952,9 @@ def forward_inference(
         **kwargs,
     ) -> BaseNavitOutputWithPast:
 
+        relay_stream_und = bool(kwargs.pop("relay_stream_und", getattr(self, "_relay_stream_und", False)))
+        relay_offload_device = kwargs.pop("relay_offload_device", getattr(self, "_relay_offload_device", torch.device("cpu")))
+
         if self.apply_qwen_2_5_vl_pos_emb:
             packed_query_position_embeddings = self.rotary_emb(packed_query_sequence.unsqueeze(0), packed_query_position_ids)
             kwargs.update({"apply_qwen_2_5_vl_pos_emb": self.apply_qwen_2_5_vl_pos_emb})
@@ -933,6 +969,11 @@ def forward_inference(
         extra_inputs = {}
         if self.use_moe:
             extra_inputs.update(mode=mode)
+            if relay_stream_und:
+                extra_inputs.update(
+                    relay_stream_und=relay_stream_und,
+                    relay_offload_device=relay_offload_device,
+                )
             if mode == "gen":
                 assert packed_vae_token_indexes is not None
                 assert packed_text_indexes is not None
@@ -960,10 +1001,21 @@ def forward_inference(
             if mode == "und":
                 packed_query_sequence = self.norm(packed_query_sequence)
             elif mode == "gen":
+                relay_stream_final_norm = (
+                    relay_stream_und
+                    and packed_text_indexes is not None
+                    and packed_text_indexes.numel() > 0
+                )
+                if relay_stream_final_norm:
+                    self.norm.to(device=packed_query_sequence.device, dtype=packed_query_sequence.dtype)
                 packed_query_sequence_ = torch.zeros_like(packed_query_sequence)
                 packed_query_sequence_[packed_text_indexes] = self.norm(packed_query_sequence[packed_text_indexes])
                 packed_query_sequence_[packed_vae_token_indexes] = self.norm_moe_gen(packed_query_sequence[packed_vae_token_indexes])
                 packed_query_sequence = packed_query_sequence_
+                if relay_stream_final_norm:
+                    self.norm.to(device=relay_offload_device)
+                    if torch.cuda.is_available():
+                        torch.cuda.empty_cache()
         else:
             packed_query_sequence = self.norm(packed_query_sequence)