diff --git a/.gitignore b/.gitignore index 6feaa2f..513458e 100644 --- a/.gitignore +++ b/.gitignore @@ -32,6 +32,7 @@ results/ downloads/ temps/ datasets/ +pr-evidence/ # outputs/ diff --git a/README.md b/README.md index 34ed3f9..5694780 100644 --- a/README.md +++ b/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`. | diff --git a/config/config_factory.py b/config/config_factory.py index 5797364..a23c7ce 100644 --- a/config/config_factory.py +++ b/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 diff --git a/inference_lance.py b/inference_lance.py index 27f49c4..9cefaf3 100644 --- a/inference_lance.py +++ b/inference_lance.py @@ -30,7 +30,9 @@ from torch.utils.data import DataLoader from transformers import HfArgumentParser, set_seed from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLVisionConfig +from safetensors import safe_open from safetensors.torch import load_file +from torch.nn.modules.module import _IncompatibleKeys from data.dataset_base import DataConfig, simple_custom_collate from data.data_utils import add_special_tokens @@ -76,6 +78,12 @@ TASK_X2T_IMAGE, TASK_X2T_VIDEO, } +MEMORY_MODE_PARALLEL = "parallel" +MEMORY_MODE_RELAY = "relay" +VALID_MEMORY_MODES = { + MEMORY_MODE_PARALLEL, + MEMORY_MODE_RELAY, +} TASK_DEFAULT_CONFIGS = { TASK_T2I: { "model_family": "image", @@ -114,7 +122,95 @@ }, } -def init_from_model_path_if_needed(model: Qwen2ForCausalLM, model_args: ModelArguments): +def normalize_memory_mode(memory_mode: str) -> str: + memory_mode = (memory_mode or MEMORY_MODE_PARALLEL).strip().lower() + if memory_mode not in VALID_MEMORY_MODES: + raise ValueError(f"memory_mode must be one of {sorted(VALID_MEMORY_MODES)}, got {memory_mode!r}") + return memory_mode + + +def validate_memory_mode(memory_mode: str, inference_args: InferenceArguments) -> None: + if memory_mode == MEMORY_MODE_RELAY: + if inference_args.task != TASK_T2I: + raise ValueError("memory_mode='relay' is currently implemented only for t2i image generation.") + if not inference_args.use_KVcache: + raise ValueError("memory_mode='relay' requires --use_KVcache true.") + + +def log_cuda_memory(stage_name: str, enabled: bool, rank0_log=print, device: Optional[int] = None) -> None: + if not enabled 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) + rank0_log(f"[memory] {stage_name}: allocated={allocated:.2f}GiB reserved={reserved:.2f}GiB peak={peak:.2f}GiB") + + +def wan_vae_config() -> AutoEncoderParams: + return AutoEncoderParams( + downsample_spatial=16, + downsample_temporal=4, + z_channels=48, + ) + + +def set_module_tensor_by_name(root: torch.nn.Module, key: str, tensor: torch.Tensor) -> bool: + parts = key.split(".") + module = root + for part in parts[:-1]: + module = getattr(module, part, None) + if module is None: + return False + + leaf_name = parts[-1] + if leaf_name in module._parameters: + old_param = module._parameters[leaf_name] + if old_param is None: + return False + if old_param.device.type == "meta": + if tensor.dtype != old_param.dtype: + tensor = tensor.to(dtype=old_param.dtype) + module._parameters[leaf_name] = torch.nn.Parameter(tensor, requires_grad=old_param.requires_grad) + return True + if tensor.dtype != old_param.dtype or tensor.device != old_param.device: + tensor = tensor.to(device=old_param.device, dtype=old_param.dtype) + old_param.data = tensor + return True + + if leaf_name in module._buffers: + old_buffer = module._buffers[leaf_name] + if old_buffer is not None and old_buffer.device.type == "meta": + if tensor.dtype != old_buffer.dtype: + tensor = tensor.to(dtype=old_buffer.dtype) + module._buffers[leaf_name] = tensor + return True + if old_buffer is not None and (tensor.dtype != old_buffer.dtype or tensor.device != old_buffer.device): + tensor = tensor.to(device=old_buffer.device, dtype=old_buffer.dtype) + module._buffers[leaf_name] = tensor + return True + + return False + + +def materialize_rotary_embedding(rotary_emb: torch.nn.Module) -> None: + inv_freq = getattr(rotary_emb, "inv_freq", None) + if inv_freq is None or inv_freq.device.type != "meta": + return + + rope_kwargs = getattr(rotary_emb, "rope_kwargs", {}) + try: + inv_freq, attention_scaling = rotary_emb.rope_init_fn(rotary_emb.config, "cpu", **rope_kwargs) + except TypeError: + inv_freq, attention_scaling = rotary_emb.rope_init_fn(rotary_emb.config, "cpu") + rotary_emb.attention_scaling = attention_scaling + rotary_emb.register_buffer("inv_freq", inv_freq, persistent=False) + rotary_emb.original_inv_freq = rotary_emb.inv_freq + + +def init_from_model_path_if_needed(model: Qwen2ForCausalLM, model_args: ModelArguments, stream: bool = False, progress_log=None): # Always load the trained Lance checkpoint from model_path. path_dir = model_args.model_path ema_path = osp.join(path_dir, "ema.safetensors") @@ -134,6 +230,47 @@ def init_from_model_path_if_needed(model: Qwen2ForCausalLM, model_args: ModelArg f"Fine-tuning failed: No valid checkpoint ('ema.safetensors' or 'model.safetensors') found in {path_dir}" ) + if stream: + expected_meta = { + name: tuple(param.shape) + for name, param in model.named_parameters() + } + expected_meta.update( + { + name: tuple(buffer.shape) + for name, buffer in model.named_buffers() + if buffer is not None + } + ) + expected_keys = set(expected_meta.keys()) + loaded_keys = set() + unexpected_keys = [] + + with safe_open(model_path_ft, framework="pt", device="cpu") as checkpoint: + checkpoint_keys = list(checkpoint.keys()) + for idx, key in enumerate(checkpoint_keys, start=1): + if key == "latent_pos_embed.pos_embed": + continue + if key not in expected_keys: + unexpected_keys.append(key) + continue + + tensor = checkpoint.get_tensor(key) + target_shape = expected_meta[key] + if target_shape != tuple(tensor.shape): + raise RuntimeError( + f"Checkpoint shape mismatch for {key}: expected {target_shape}, got {tuple(tensor.shape)}" + ) + if not set_module_tensor_by_name(model, key, tensor): + unexpected_keys.append(key) + continue + loaded_keys.add(key) + if progress_log is not None and len(loaded_keys) % 100 == 0: + progress_log(f"[startup] streamed {len(loaded_keys)}/{len(checkpoint_keys)} checkpoint tensors") + + missing_keys = [key for key in expected_keys if key not in loaded_keys] + return _IncompatibleKeys(missing_keys, unexpected_keys) + # NOTE: position embeds are fixed sinusoidal embeddings, so we can just pop it off, # which makes it easier to adapt to different resolutions. if 'latent_pos_embed.pos_embed' in model_state_dict: @@ -264,13 +401,34 @@ def validate_on_fixed_batch( save_path_gen: str = "", save_path_gt: str = "", ): + memory_mode = normalize_memory_mode(inference_args.memory_mode) + relay_enabled = memory_mode == MEMORY_MODE_RELAY + + def ensure_vae_model(current_vae_model: Optional[WanVideoVAE]) -> WanVideoVAE: + if current_vae_model is None: + current_vae_model = WanVideoVAE(device="cpu") + return current_vae_model + val_data = val_data_cpu.cuda(device).to_dict() - fsdp_model = fsdp_model.to(device=device, dtype=torch.bfloat16) + if relay_enabled: + fsdp_model.configure_relay_memory(enabled=True, offload_device="cpu", log_memory=inference_args.relay_memory_log) + else: + fsdp_model = fsdp_model.to(device=device, dtype=torch.bfloat16) + log_cuda_memory("validation batch ready", inference_args.relay_memory_log, device=device) with torch.no_grad(), torch.amp.autocast("cuda", enabled=True, dtype=torch.bfloat16): # Compute padded_latent. if "padded_videos" in val_data.keys(): + needs_vae_encode = any(str(mode).lower().startswith("on") for mode in val_data["vae_data_mode"]) + if relay_enabled and needs_vae_encode: + vae_model = ensure_vae_model(vae_model) + vae_model.to(torch.device("cuda", device)) + log_cuda_memory("after VAE materialize for encode", inference_args.relay_memory_log, device=device) val_data["padded_latent"] = make_padded_latent(val_data["padded_videos"], val_data["vae_data_mode"], vae_model) + if relay_enabled and needs_vae_encode and vae_model is not None: + vae_model.to("cpu") + clean_memory() + log_cuda_memory("after VAE encode offload", inference_args.relay_memory_log, device=device) # -------------------- Generation branch -------------------- if inference_args.task in GENERATION_TASKS: @@ -313,12 +471,23 @@ def validate_on_fixed_batch( "cfg_uncond_token_id": training_args.cfg_uncond_token_id, "index": val_data["index"], "val_padded_videos": val_data["padded_videos"] if save_source_video else None, + "memory_mode": memory_mode, + "relay_memory_log": inference_args.relay_memory_log, } if inference_args.use_KVcache: denoise_latent, captions, padded_videos, index = fsdp_model.validation_gen_KVcache(**params) else: denoise_latent, captions, padded_videos, index = fsdp_model.validation_gen(**params) + if relay_enabled: + log_cuda_memory("before Lance offload for VAE decode", inference_args.relay_memory_log, device=device) + fsdp_model.relay_discard_all() + clean_memory() + log_cuda_memory("after Lance release for VAE decode", inference_args.relay_memory_log, device=device) + vae_model = ensure_vae_model(vae_model) + vae_model.to(torch.device("cuda", device)) + log_cuda_memory("after VAE materialize", inference_args.relay_memory_log, device=device) + # Decode. for i_val, latent in enumerate(denoise_latent): if inference_args.task in {TASK_I2V, TASK_IMAGE_EDIT, TASK_VIDEO_EDIT}: @@ -348,6 +517,10 @@ def validate_on_fixed_batch( clean_memory() del denoise_latent, captions, padded_videos, params + if relay_enabled and vae_model is not None: + vae_model.to("cpu") + clean_memory() + log_cuda_memory("after VAE offload", inference_args.relay_memory_log, device=device) clean_memory() elif inference_args.task in UNDERSTANDING_TASKS: @@ -421,10 +594,14 @@ def main(): # ========================= Load task paths and example JSONs from defaults ============================== apply_inference_defaults(model_args, data_args, inference_args) + inference_args.memory_mode = normalize_memory_mode(inference_args.memory_mode) + validate_memory_mode(inference_args.memory_mode, inference_args) training_args.validation_noise_seed = training_args.validation_data_seed + relay_enabled = inference_args.memory_mode == MEMORY_MODE_RELAY logger = get_logger() - log_rank0 = print if GLOBAL_RANK == 0 else (lambda *_: None) # Only print on rank 0. + log_rank0 = (lambda *args, **kwargs: print(*args, **kwargs, flush=True)) if GLOBAL_RANK == 0 else (lambda *_args, **_kwargs: None) # Only print on rank 0. + log_rank0(f"[startup] memory_mode={inference_args.memory_mode}") def log_stage(stage_name: str, start_time: float, extra: str = ""): elapsed = time.perf_counter() - start_time @@ -436,79 +613,106 @@ def log_stage(stage_name: str, start_time: float, extra: str = ""): set_seed(seed) # ========================= LLM model setup ============================== - stage_start = time.perf_counter() - log_rank0(f"[startup] Loading LLM config: {osp.join(model_args.model_path, 'llm_config.json')}") - llm_config: Qwen2Config = Qwen2Config.from_json_file(osp.join(model_args.model_path, "llm_config.json")) - log_stage("LLM config load", stage_start) + relay_t2i_without_vit = relay_enabled and inference_args.task == TASK_T2I + load_visual_und = training_args.visual_und and not relay_t2i_without_vit + vit_config = None + vit_model = None - llm_config.layer_module = model_args.layer_module - llm_config.qk_norm = model_args.llm_qk_norm - llm_config.qk_norm_und = model_args.llm_qk_norm_und - llm_config.qk_norm_gen = model_args.llm_qk_norm_gen + original_default_dtype = torch.get_default_dtype() + if relay_enabled: + torch.set_default_dtype(torch.bfloat16) + log_rank0("[startup] Using bf16 default dtype for relay CPU module construction") - llm_config.tie_word_embeddings = model_args.tie_word_embeddings - llm_config.freeze_und = training_args.freeze_und - llm_config.apply_qwen_2_5_vl_pos_emb = training_args.apply_qwen_2_5_vl_pos_emb + try: + stage_start = time.perf_counter() + log_rank0(f"[startup] Loading LLM config: {osp.join(model_args.model_path, 'llm_config.json')}") + llm_config: Qwen2Config = Qwen2Config.from_json_file(osp.join(model_args.model_path, "llm_config.json")) + log_stage("LLM config load", stage_start) - stage_start = time.perf_counter() - log_rank0(f"[startup] Initializing LLM weights: {model_args.model_path}") - language_model: Qwen2ForCausalLM = Qwen2ForCausalLM(llm_config) - log_stage("LLM weight init", stage_start) + llm_config.layer_module = model_args.layer_module + llm_config.qk_norm = model_args.llm_qk_norm + llm_config.qk_norm_und = model_args.llm_qk_norm_und + llm_config.qk_norm_gen = model_args.llm_qk_norm_gen - if training_args.visual_und: - if model_args.vit_type in ("qwen2_5_vl", "qwen_2_5_vl_original"): - stage_start = time.perf_counter() - log_rank0(f"[startup] Loading VIT config: {model_args.vit_path}") - vit_config = Qwen2_5_VLVisionConfig.from_pretrained(model_args.vit_path) - log_stage("VIT config load", stage_start) + llm_config.tie_word_embeddings = model_args.tie_word_embeddings + llm_config.freeze_und = training_args.freeze_und + llm_config.apply_qwen_2_5_vl_pos_emb = training_args.apply_qwen_2_5_vl_pos_emb - stage_start = time.perf_counter() - log_rank0(f"[startup] Loading VIT weights: {osp.join(model_args.vit_path, 'vit.safetensors')}") - vit_model = Qwen2_5_VisionTransformerPretrainedModel(vit_config) - vit_weights = load_file(osp.join(model_args.vit_path, "vit.safetensors")) - vit_model.load_state_dict(vit_weights, strict=True) - log_stage("VIT weight load", stage_start) + stage_start = time.perf_counter() + log_rank0(f"[startup] Initializing LLM weights: {model_args.model_path}") + if relay_enabled: + with torch.device("meta"): + language_model: Qwen2ForCausalLM = Qwen2ForCausalLM(llm_config) else: - raise ValueError(f"Unsupported vit_type: {model_args.vit_type}") - - clean_memory(vit_weights) + language_model: Qwen2ForCausalLM = Qwen2ForCausalLM(llm_config) + log_stage("LLM weight init", stage_start) + + if load_visual_und: + if model_args.vit_type in ("qwen2_5_vl", "qwen_2_5_vl_original"): + stage_start = time.perf_counter() + log_rank0(f"[startup] Loading VIT config: {model_args.vit_path}") + vit_config = Qwen2_5_VLVisionConfig.from_pretrained(model_args.vit_path) + log_stage("VIT config load", stage_start) + + stage_start = time.perf_counter() + log_rank0(f"[startup] Loading VIT weights: {osp.join(model_args.vit_path, 'vit.safetensors')}") + vit_model = Qwen2_5_VisionTransformerPretrainedModel(vit_config) + vit_weights = load_file(osp.join(model_args.vit_path, "vit.safetensors")) + vit_model.load_state_dict(vit_weights, strict=True) + log_stage("VIT weight load", stage_start) + clean_memory(vit_weights) + else: + raise ValueError(f"Unsupported vit_type: {model_args.vit_type}") + elif training_args.visual_und: + log_rank0("[startup] Skipping VIT weight load for t2i relay") - if training_args.visual_gen: - stage_start = time.perf_counter() - log_rank0("[startup] Initializing VAE") - vae_model = WanVideoVAE() - vae_config: AutoEncoderParams = deepcopy(vae_model.vae_config) - log_stage("VAE init", stage_start) + if training_args.visual_gen: + stage_start = time.perf_counter() + if relay_enabled: + vae_model = None + vae_config = wan_vae_config() + log_rank0("[startup] Deferring VAE weight load until relay decode") + else: + log_rank0(f"[startup] Initializing VAE on {DEVICE}") + vae_model = WanVideoVAE(device=DEVICE) + vae_config: AutoEncoderParams = deepcopy(vae_model.vae_config) + log_stage("VAE init", stage_start) + else: + vae_model = None + vae_config = None + + # Lance configuration + config = LanceConfig( + visual_gen=training_args.visual_gen, + visual_und=load_visual_und, + llm_config=llm_config, + vit_config=vit_config if load_visual_und else None, + vae_config=vae_config if training_args.visual_gen else None, + latent_patch_size=model_args.latent_patch_size, + max_num_frames=model_args.max_num_frames, + max_latent_size=model_args.max_latent_size, + vit_max_num_patch_per_side=model_args.vit_max_num_patch_per_side, + connector_act=model_args.connector_act, + interpolate_pos=model_args.interpolate_pos, + timestep_shift=training_args.timestep_shift, + ) + model: Lance = Lance( + language_model=language_model, + vit_model=vit_model if load_visual_und else None, + vit_type=model_args.vit_type, + config=config, + training_args=training_args, + ) + finally: + if relay_enabled: + torch.set_default_dtype(original_default_dtype) + if relay_enabled: + log_rank0("[startup] Keeping Lance model on CPU for relay phase loading") else: - vae_model = None - vae_config = None - - # Lance configuration - config = LanceConfig( - visual_gen=training_args.visual_gen, - visual_und=training_args.visual_und, - llm_config=llm_config, - vit_config=vit_config if training_args.visual_und else None, - vae_config=vae_config if training_args.visual_gen else None, - latent_patch_size=model_args.latent_patch_size, - max_num_frames=model_args.max_num_frames, - max_latent_size=model_args.max_latent_size, - vit_max_num_patch_per_side=model_args.vit_max_num_patch_per_side, - connector_act=model_args.connector_act, - interpolate_pos=model_args.interpolate_pos, - timestep_shift=training_args.timestep_shift, - ) - model: Lance = Lance( - language_model=language_model, - vit_model=vit_model if training_args.visual_und else None, - vit_type=model_args.vit_type, - config=config, - training_args=training_args, - ) - stage_start = time.perf_counter() - log_rank0(f"[startup] Moving Lance model to GPU {DEVICE}") - model = model.to(DEVICE) - log_stage("Lance model move to GPU", stage_start) + stage_start = time.perf_counter() + log_rank0(f"[startup] Moving Lance model to GPU {DEVICE}") + model = model.to(DEVICE) + log_stage("Lance model move to GPU", stage_start) # Setup tokenizer for model: stage_start = time.perf_counter() @@ -519,10 +723,16 @@ def log_stage(stage_name: str, start_time: float, extra: str = ""): log_stage("tokenizer load and special token init", stage_start, extra=f"num_new_tokens={num_new_tokens}") # Initialize MoE before loading the checkpoint. - if training_args.copy_init_moe: + if training_args.copy_init_moe and not relay_enabled: language_model.init_moe() + elif training_args.copy_init_moe and relay_enabled: + log_rank0("[startup] Skipping pre-load MoE copy for relay streaming checkpoint load") - init_from_model_path_if_needed(model, model_args) + if relay_enabled: + log_rank0("[startup] Loading Lance checkpoint with streaming CPU loader") + init_from_model_path_if_needed(model, model_args, stream=relay_enabled, progress_log=log_rank0 if relay_enabled else None) + if relay_enabled and hasattr(model.language_model.model, "rotary_emb"): + materialize_rotary_embedding(model.language_model.model.rotary_emb) # Resize afterward to avoid checkpoint shape mismatches or overwritten weights. if num_new_tokens > 0: @@ -549,7 +759,11 @@ def log_stage(stage_name: str, start_time: float, extra: str = ""): else: assert model.language_model.get_input_embeddings().weight.data.data_ptr() != model.language_model.get_output_embeddings().weight.data.data_ptr(), 'tie_word_embeddings conflict' - model = model.to(device=DEVICE, dtype=torch.bfloat16) + if relay_enabled: + model = model.to(dtype=torch.bfloat16) + model.configure_relay_memory(enabled=True, offload_device="cpu", log_memory=inference_args.relay_memory_log) + else: + model = model.to(device=DEVICE, dtype=torch.bfloat16) model.eval() if vae_model is not None and hasattr(vae_model, "eval"): vae_model.eval() diff --git a/inference_lance.sh b/inference_lance.sh index 3a5b959..b5504d9 100755 --- a/inference_lance.sh +++ b/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 "" diff --git a/modeling/lance/lance.py b/modeling/lance/lance.py index c425fe6..ff63c29 100644 --- a/modeling/lance/lance.py +++ b/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) diff --git a/modeling/lance/qwen2_navit.py b/modeling/lance/qwen2_navit.py index 8f3c8d8..ef4abc2 100644 --- a/modeling/lance/qwen2_navit.py +++ b/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) diff --git a/tools/compare_memory_modes_t2i.py b/tools/compare_memory_modes_t2i.py new file mode 100644 index 0000000..f46f7ce --- /dev/null +++ b/tools/compare_memory_modes_t2i.py @@ -0,0 +1,246 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Compare text-to-image outputs from parallel and relay memory modes.""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +import json +import shlex +import subprocess +import sys +import time +from pathlib import Path + + +RESOLUTION_SIZES = { + "image_256res": (256, 256), + "image_512res": (512, 512), + "image_768res": (768, 768), +} +DEFAULT_PROMPT = 'A cat holds a poster with rainbow text "STOP"' +MODES = ("parallel", "relay") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-path", default="downloads/Lance_3B", help="Path to Lance T2I model weights.") + parser.add_argument( + "--resolutions", + nargs="+", + default=["image_256res"], + choices=sorted(RESOLUTION_SIZES), + help="T2I resolution presets to compare.", + ) + parser.add_argument("--prompt", default=DEFAULT_PROMPT, help="Prompt used for both memory modes.") + parser.add_argument("--output-root", default="results/memory_mode_compare_t2i", help="Where outputs and logs are written.") + parser.add_argument("--timesteps", type=int, default=30, help="Validation denoising steps.") + parser.add_argument("--timestep-shift", type=float, default=3.5, help="Validation timestep shift.") + parser.add_argument("--seed", type=int, default=42, help="Validation data seed.") + parser.add_argument("--cfg-text-scale", type=float, default=4.0, help="Text CFG scale.") + parser.add_argument("--shell", default="bash", help="Shell executable used to run inference_lance.sh.") + parser.add_argument("--relay-memory-log", action="store_true", help="Forward --RELAY_MEMORY_LOG true to inference.") + parser.add_argument("--reuse-existing", action="store_true", help="Hash existing outputs instead of rerunning finished modes.") + parser.add_argument("--dry-run", action="store_true", help="Print commands without running inference.") + return parser.parse_args() + + +def repo_root() -> Path: + return Path(__file__).resolve().parents[1] + + +def sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def write_prompt_config(output_root: Path, prompt: str) -> Path: + config_dir = output_root / "configs" + config_dir.mkdir(parents=True, exist_ok=True) + config_path = config_dir / "t2i_prompt.json" + with config_path.open("w", encoding="utf-8") as handle: + json.dump({"000000.png": prompt}, handle, ensure_ascii=False, indent=2) + handle.write("\n") + return config_path + + +def output_png(save_dir: Path) -> Path: + candidate = save_dir / "000000.png" + if candidate.exists(): + return candidate + pngs = sorted(save_dir.glob("*.png")) + if pngs: + return pngs[0] + raise FileNotFoundError(f"No PNG output found in {save_dir}") + + +def build_command(args: argparse.Namespace, root: Path, config_path: Path, resolution: str, mode: str, save_dir: Path) -> list[str]: + height, width = RESOLUTION_SIZES[resolution] + return [ + args.shell, + str(root / "inference_lance.sh"), + "--TASK_NAME", + "t2i", + "--MODEL_PATH", + args.model_path, + "--CONFIG_PATH", + str(config_path), + "--RESOLUTION", + resolution, + "--VIDEO_HEIGHT", + str(height), + "--VIDEO_WIDTH", + str(width), + "--NUM_FRAMES", + "1", + "--VALIDATION_NUM_TIMESTEPS", + str(args.timesteps), + "--VALIDATION_TIMESTEP_SHIFT", + str(args.timestep_shift), + "--VALIDATION_DATA_SEED", + str(args.seed), + "--CFG_TEXT_SCALE", + str(args.cfg_text_scale), + "--USE_KVCACHE", + "true", + "--MEMORY_MODE", + mode, + "--RELAY_MEMORY_LOG", + str(args.relay_memory_log).lower(), + "--SAVE_PATH_GEN", + str(save_dir), + ] + + +def run_mode(args: argparse.Namespace, root: Path, config_path: Path, output_root: Path, resolution: str, mode: str) -> dict: + save_dir = output_root / f"{resolution}_{mode}" + log_dir = output_root / "logs" + log_dir.mkdir(parents=True, exist_ok=True) + log_path = log_dir / f"{resolution}_{mode}.log" + cmd = build_command(args, root, config_path, resolution, mode, save_dir) + printable_cmd = " ".join(shlex.quote(part) for part in cmd) + + if args.dry_run: + print(printable_cmd) + return { + "resolution": resolution, + "mode": mode, + "exit_code": "", + "seconds": "", + "sha256": "", + "path": str(save_dir / "000000.png"), + "log": str(log_path), + } + + start = time.perf_counter() + if args.reuse_existing and (save_dir / "000000.png").exists(): + exit_code = 0 + else: + save_dir.mkdir(parents=True, exist_ok=True) + with log_path.open("w", encoding="utf-8") as log_file: + log_file.write(f"$ {printable_cmd}\n\n") + log_file.flush() + result = subprocess.run(cmd, cwd=root, stdout=log_file, stderr=subprocess.STDOUT, text=True, check=False) + exit_code = result.returncode + + seconds = time.perf_counter() - start + image_path = output_png(save_dir) if exit_code == 0 else save_dir / "000000.png" + image_hash = sha256_file(image_path) if image_path.exists() else "" + + return { + "resolution": resolution, + "mode": mode, + "exit_code": str(exit_code), + "seconds": f"{seconds:.2f}", + "sha256": image_hash, + "path": str(image_path), + "log": str(log_path), + } + + +def write_summary(output_root: Path, rows: list[dict]) -> None: + csv_path = output_root / "summary.csv" + fieldnames = ["resolution", "mode", "exit_code", "seconds", "sha256", "path", "log"] + with csv_path.open("w", encoding="utf-8", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(rows) + + md_path = output_root / "summary.md" + with md_path.open("w", encoding="utf-8") as handle: + handle.write("| Resolution | Mode | Exit code | Seconds | SHA256 |\n") + handle.write("| --- | --- | ---: | ---: | --- |\n") + for row in rows: + handle.write( + f"| {row['resolution']} | {row['mode']} | {row['exit_code']} | {row['seconds']} | {row['sha256']} |\n" + ) + + +def compare_rows(rows: list[dict]) -> int: + exit_code = 0 + by_resolution: dict[str, dict[str, dict]] = {} + for row in rows: + by_resolution.setdefault(row["resolution"], {})[row["mode"]] = row + + for resolution, modes in by_resolution.items(): + parallel = modes.get("parallel") + relay = modes.get("relay") + if parallel is None or relay is None: + print(f"{resolution}: missing comparison row", file=sys.stderr) + exit_code = 2 + continue + if parallel["exit_code"] != "0" or relay["exit_code"] != "0": + print(f"{resolution}: inference failed; see logs", file=sys.stderr) + exit_code = 1 + continue + if not parallel["sha256"] or not relay["sha256"]: + print(f"{resolution}: missing output hash", file=sys.stderr) + exit_code = 2 + continue + if parallel["sha256"] != relay["sha256"]: + print(f"{resolution}: hash mismatch", file=sys.stderr) + exit_code = 2 + continue + print(f"{resolution}: bit-identical") + + return exit_code + + +def main() -> int: + args = parse_args() + root = repo_root() + output_root = (root / args.output_root).resolve() + output_root.mkdir(parents=True, exist_ok=True) + config_path = write_prompt_config(output_root, args.prompt) + + rows = [] + for resolution in args.resolutions: + for mode in MODES: + rows.append(run_mode(args, root, config_path, output_root, resolution, mode)) + + write_summary(output_root, rows) + if args.dry_run: + print(f"Dry run complete. Summary scaffold: {output_root}") + return 0 + return compare_rows(rows) + + +if __name__ == "__main__": + raise SystemExit(main())