EthanZyh's picture
copied from EthanZyh/DiffusionText2WorldGeneration
8c31d70
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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.
import json
import os
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Set
from .log import log
import torch
from safetensors.torch import load_file
from torch.nn.modules.module import _IncompatibleKeys
from .ar_configs_base_model import ModelConfig
from .ar_config_tokenizer import TokenizerConfig
from .mm_projector import MultimodalProjector
from .ar_transformer import Transformer
from .vit import VisionTransformer, get_vit_config
from .ar_tokenizer import DiscreteMultimodalTokenizer, update_vocab_size
from .checkpoint import (
get_partial_state_dict,
process_state_dict,
substrings_to_ignore,
)
from .sampling import decode_n_tokens, decode_one_token, prefill
from .misc import misc, Color, timer
class AutoRegressiveModel(torch.nn.Module):
"""
A class to build and use a AutoRegressiveModel model for text generation.
Methods:
build: Build a AutoRegressiveModel instance by initializing and loading a model checkpoint.
generate: Generate text sequences based on provided prompts using the language generation model.
"""
def __init__(
self,
model: Transformer = None,
tokenizer: DiscreteMultimodalTokenizer = None,
config: ModelConfig = None,
vision_encoder: VisionTransformer = None,
mm_projector: MultimodalProjector = None,
):
"""
Initialize the AutoRegressiveModel instance with a model and tokenizer.
Args:
model (Transformer): The Transformer model for text generation.
tokenizer (Tokenizer): The tokenizer for encoding and decoding text.
config (Config): The configuration for the AutoRegressiveModel model.
vision_encoder (VisionTransformer): The vision encoder for the AutoRegressiveModel model.
mm_projector (MultimodalProjector): The multi-modal projector for the AutoRegressiveModel model.
"""
super().__init__()
self.model = model
self.tokenizer = tokenizer
self.config = config
self.vision_encoder = vision_encoder
self.mm_projector = mm_projector
@property
def precision(self):
return self.model.precision
def get_num_params(
self,
) -> int:
"""
Return the number of parameters in the model.
"""
n_params = sum(p.numel() for p in self.parameters())
return n_params
def load_ar_model(
self,
tokenizer_config,
):
"""
Load the AR model.
"""
model_config = self.config
ckpt_path = model_config.ckpt_path
with timer(f"loading checkpoint from {ckpt_path}"):
if ckpt_path.endswith("safetensors"):
# Load with safetensors API
checkpoint = load_file(ckpt_path, device="cpu")
else:
# The pytorch version
checkpoint = torch.load(
ckpt_path,
map_location="cpu",
mmap=True, # load the checkpoint in memory-mapped mode
weights_only=True,
)
llm_checkpoint = checkpoint["model"] if "model" in checkpoint else checkpoint
orig_precision = torch.get_default_dtype()
precision = getattr(torch, model_config.precision)
torch.set_default_dtype(precision)
log.debug(f"Setting torch default dtype to {precision}")
model = Transformer(
params=model_config,
tokenizer_config=tokenizer_config,
)
log.debug(
f"tokenizer tokenizer_config.video_tokenizer.vocab_size {tokenizer_config.video_tokenizer.vocab_size}"
)
vocab_size = update_vocab_size(
existing_vocab_size=0,
to_be_added_vocab_size=tokenizer_config.video_tokenizer.vocab_size,
training_type=tokenizer_config.training_type,
add_special_tokens=False,
)
log.debug(
f"tokenizer tokenizer_config.video_tokenizer.vocab_size {tokenizer_config.video_tokenizer.vocab_size} vocab_size {vocab_size}"
)
# Perform vocab expansion
if vocab_size > model.vocab_size:
log.debug(f"Expanding vocab size to {vocab_size}")
# For text-to-video training, we only expand the embedding layer but not the output (unembedding) layer,
expand_output_layer = not (tokenizer_config.training_type == "text_to_video")
model.expand_vocab(
vocab_size,
init_method="gaussian",
expand_output_layer=expand_output_layer,
)
# Remove the "model." prefix in the state_dict
llm_checkpoint = process_state_dict(llm_checkpoint, prefix_to_remove="model.")
with timer("loading state_dict into model"):
missing_keys, _ = model.load_state_dict(llm_checkpoint, strict=True)
# Remove keys with "_extra_state" suffix in missing_keys (defined by TransformerEngine for FP8 usage)
missing_keys = [k for k in missing_keys if not k.endswith("_extra_state")]
assert len(missing_keys) == 0, f"Missing keys: {missing_keys}"
self.model = model.to(precision).to("cuda")
torch.set_default_dtype(orig_precision) # Reset the default dtype to the original value
def load_tokenizer(self, tokenizer_config):
"""
Load the tokenizer.
"""
self.tokenizer = DiscreteMultimodalTokenizer(tokenizer_config)
@staticmethod
def build(
model_config: ModelConfig = ModelConfig(),
tokenizer_config: TokenizerConfig = None,
) -> "AutoRegressiveModel":
"""
Build a AutoRegressiveModel instance by initializing and loading a model checkpoint.
Args:
model_config (ModelConfig, optional): The model configuration for the AutoRegressiveModel instance. Defaults to ModelConfig().
tokenizer_config (TokenizerConfig, optional): The tokenizer configuration for the AutoRegressiveModel instance. Defaults to None.
download_rank_sync (bool, optional): Whether to download the checkpoint in a rank-synchronized manner. Defaults to True.
Returns:
AutoRegressiveModel: An instance of the AutoRegressiveModel class with the loaded model and tokenizer.
Raises:
AssertionError: If there are no checkpoint files in the specified directory.
Note:
This method sets the device to CUDA and loads the pre-trained model and tokenizer.
"""
# Initialize model configuration parameters
config_params = {}
# Load checkpoint and model parameters
if model_config.ckpt_path is None:
# If ckpt_path is not provided, we assume the model checkpoint is saved in the ckpt_dir
ckpt_dir = model_config.ckpt_dir
# We prioritize safetensors version over the pytorch version, since the former is
# much faster for checkpoint loading.
checkpoints = sorted(Path(ckpt_dir).glob("*.safetensors"))
if len(checkpoints) == 0:
checkpoints = sorted(Path(ckpt_dir).glob("*.pth"))
assert len(checkpoints) > 0, f"no checkpoint files found in {ckpt_dir}"
assert (
len(checkpoints) == 1
), f"multiple checkpoint files found in {ckpt_dir} (currently only one is supported)"
ckpt_path = str(checkpoints[0]) # Assuming single checkpoint for non-parallel case
if os.path.exists(Path(ckpt_dir) / "config.json"):
with open(Path(ckpt_dir) / "config.json", "r") as f:
config_params = json.loads(f.read())
else:
log.info(
f"No params.json found in the checkpoint directory ({ckpt_dir}). " f"Using default model config."
)
else:
# If ckpt_path is provided, we load the model from the specified path,
# and use the default model configuration
ckpt_path = model_config.ckpt_path
for key, value in config_params.items():
if hasattr(model_config, key):
# Override the default model configuration with the parameters from the checkpoint
setattr(model_config, key, value)
with timer(f"loading checkpoint from {ckpt_path}"):
if ckpt_path.endswith("safetensors"):
# Load with safetensors API
checkpoint = load_file(ckpt_path, device="cpu")
else:
# The pytorch version
checkpoint = torch.load(
ckpt_path,
map_location="cpu",
mmap=True, # load the checkpoint in memory-mapped mode
weights_only=True,
)
llm_checkpoint = checkpoint["model"] if "model" in checkpoint else checkpoint
if model_config.vision_encoder is not None:
# Take the LLM weights (starting with "model.") from the VLM checkpoint
llm_checkpoint = get_partial_state_dict(llm_checkpoint, prefix="model.")
if model_config.vision_encoder is not None:
# For vanilla VLM ckpt before fine-tuning, `checkpoint['model']` only contains LLM weights, and `checkpoint['vision_encoder']`
# and `checkpoint['mm_projector']` are both for those weights
# For fine-tuned VLM ckpt, `checkpoint['model']` contains all LLM, mm_projector and vision_encoder weights
if "vision_encoder" in checkpoint:
log.debug("Using pretrained vision_encoder")
vit_checkpoint = checkpoint["vision_encoder"]
else:
log.debug("Using fine-tuned vision_encoder")
vit_checkpoint = get_partial_state_dict(llm_checkpoint, prefix="vision_encoder.")
vit_checkpoint = process_state_dict(vit_checkpoint, prefix_to_remove="vision_encoder.")
if "mm_projector" in checkpoint:
log.debug("Using pretrained mm_projector")
projector_checkpoint = checkpoint["mm_projector"]
else:
log.debug("Using fine-tuned mm_projector")
projector_checkpoint = get_partial_state_dict(llm_checkpoint, prefix="mm_projector.")
projector_checkpoint = process_state_dict(projector_checkpoint, prefix_to_remove="mm_projector.")
assert (
len(vit_checkpoint) > 0 and len(projector_checkpoint) > 0
), "vit_checkpoint and projector_checkpoint cannot be empty. We do not support random initialization for vision_encoder and mm_projector."
tokenizer = DiscreteMultimodalTokenizer(tokenizer_config)
orig_precision = torch.get_default_dtype()
precision = getattr(torch, model_config.precision)
torch.set_default_dtype(precision)
log.debug(f"Setting torch default dtype to {precision}")
model = Transformer(
params=model_config,
tokenizer_config=tokenizer_config,
)
model_kwargs = {}
if model_config.vision_encoder is not None:
assert model_config.mm_projector is not None, "mm_projector must be provided if vision_encoder is provided."
vit_config = get_vit_config(model_config.vision_encoder)
vision_encoder = VisionTransformer.build(
vit_config,
)
mm_projector = MultimodalProjector(
mm_projector_type=model_config.mm_projector, in_dim=vit_config["dim"], out_dim=model_config["dim"]
)
model_kwargs.update({"vision_encoder": vision_encoder, "mm_projector": mm_projector})
# Perform vocab expansion
if tokenizer.vocab_size > model.vocab_size:
log.debug(f"Expanding vocab size to {tokenizer.vocab_size}")
# For text-to-video training, we only expand the embedding layer but not the output (unembedding) layer,
expand_output_layer = not (tokenizer.training_type == "text_to_video")
model.expand_vocab(
tokenizer.vocab_size,
init_method="gaussian",
expand_output_layer=expand_output_layer,
)
# Remove the "model." prefix in the state_dict
llm_checkpoint = process_state_dict(llm_checkpoint, prefix_to_remove="model.")
with timer("loading state_dict into model"):
missing_keys, unexpected_keys = model.load_state_dict(llm_checkpoint, strict=True)
# Remove keys with "_extra_state" suffix in missing_keys (defined by TransformerEngine for FP8 usage)
missing_keys = [k for k in missing_keys if not k.endswith("_extra_state")]
assert len(missing_keys) == 0, f"Missing keys: {missing_keys}"
if model_config.vision_encoder is not None:
vision_encoder.load_state_dict(vit_checkpoint)
mm_projector.load_state_dict(projector_checkpoint)
if model_config.vision_encoder_in_channels != 3:
vision_encoder.expand_in_channels(model_config.vision_encoder_in_channels)
model = model.to(precision) # ensure model parameters are in the correct precision
log.debug(f"Model config: {model_config}")
model_class = AutoRegressiveModel
torch.set_default_dtype(orig_precision) # Reset the default dtype to the original value
return model_class(model, tokenizer, model_config, **model_kwargs)
@torch.no_grad()
def generate(
self,
prompt_tokens: List[List[int]] | torch.Tensor,
max_gen_len: int,
temperature: float = 1.0,
top_k: Optional[int] = None,
top_p: Optional[float] = None,
num_gen_seq: int = 1,
logprobs: bool = False,
echo: bool = False,
seed: int = None,
context: Optional[torch.Tensor] = None,
context_mask: Optional[torch.Tensor] = None,
compile_sampling: bool = True,
compile_prefill: bool = False,
verbose: bool = True,
stop_tokens: Optional[Set[int]] = None,
images: Optional[torch.Tensor] = None,
):
"""
Autoregressive generation built upon the gpt-fast implementation (https://github.com/pytorch-labs/gpt-fast).
Args:
prompt_tokens (List[List[int]] | torch.Tensor): A single prompt of shape (1, seq_len).
max_gen_len (int): Maximum length of the generated text sequence.
temperature (float, optional): Temperature value for controlling randomness in sampling. Defaults to 0.6.
top_k (int, optional): Top-k value for top-k sampling. Defaults to None.
top_p (float, optional): Top-p probability threshold for nucleus sampling. Defaults to None.
num_gen_seq (int, optional): Number of outputs to generate given the same prompt. Defaults to 1. When temperature == 0, num_gen_seq must be 1 because the generation is deterministic.
echo (bool, optional): Flag indicating whether to include prompt tokens in the generated output. Defaults to False.
logit_clipping_range (list, optional): Range of logits to clip. Defaults to [].
seed (int, optional): Random seed for reproducibility. Defaults to None.
compile_sampling (bool, optional): Flag indicating whether to compile the decoding function. Defaults to True.
compile_prefill (bool, optional): Flag indicating whether to compile the prefill function. Defaults to False.
verbose (bool, optional): Flag indicating whether to print the the time. Defaults to False.
"""
assert top_k is None or top_p is None, f"Only one of top_k ({top_k} or top_p ({top_p} should be specified."
if temperature == 0:
top_p, top_k = None, None
log.debug("Setting top_p and top_k to None because temperature is 0")
if top_p is not None:
log.debug(f"Using top-p sampling with p={top_p} and temperature={temperature}")
elif top_k is not None:
log.debug(f"Using top-k sampling with k={top_k} and temperature={temperature}")
else:
log.debug("Not applying top-k or top-p sampling. Will use top-k sampling with k=None")
orig_precision = torch.get_default_dtype()
torch.set_default_dtype(self.precision)
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.triton.unique_kernel_names = True
# Experimental features to reduce compilation times, will be on by default in future
torch._inductor.config.fx_graph_cache = True
if seed is not None:
misc.set_random_seed(seed)
assert not logprobs, "logprobs are not supported for fast_generate yet"
# Examine if the function prefil and decode_one_token functions are compiled yet. If not, compile them based on the flags
if compile_sampling and not getattr(self, "inference_decode_compiled", False):
self.decode_one_token = torch.compile(decode_one_token, mode="reduce-overhead", fullgraph=True)
self.inference_decode_compiled = True
log.info("Compiled AR sampling function. Note: the first run will be slower due to compilation")
if compile_prefill and not getattr(self, "inference_prefill_compiled", False):
self.prefill = torch.compile(prefill, fullgraph=True, dynamic=True)
self.inference_prefill_compiled = True
log.info("Compiled prefill function. Note: the first run will be slower due to compilation")
if not hasattr(self, "decode_one_token"):
self.decode_one_token = decode_one_token
if not hasattr(self, "prefill"):
self.prefill = prefill
# Initialization and Assertions
if isinstance(self.model.params, list):
# During training, model.params is a list
log.debug(
f"Find self.model.params is a list, use self.config instead. Get max_batch_size={self.config.max_batch_size}, max_seq_len={self.config.max_seq_len}"
)
params = self.config
else:
params = self.model.params
if isinstance(prompt_tokens, list):
prompt_tokens = torch.tensor(prompt_tokens, dtype=torch.long, device="cuda")
if prompt_tokens.ndim == 1:
prompt_tokens = prompt_tokens.view(1, -1)
else:
assert prompt_tokens.ndim == 2, f"prompt_tokens has shape {prompt_tokens.shape}"
batch_size, prompt_len = prompt_tokens.shape
total_len = min(params.max_seq_len, max_gen_len + prompt_len)
if max_gen_len + prompt_len > params.max_seq_len:
log.warning(
f"max_gen_len + prompt_len={max_gen_len + prompt_len} exceeds max_seq_len={params.max_seq_len}, truncate max_gen_len to {params.max_seq_len - prompt_len}"
)
max_gen_len = params.max_seq_len - prompt_len
if context_mask is not None:
context_mask = context_mask.to(dtype=torch.bool)
if context_mask.ndim == 2:
assert (
context_mask.shape[0] == batch_size
), f"batch_size mismatch: {context_mask.shape[0]} != {batch_size}"
# Unsqueeze it to make it of shape [batch_size, 1, 1, context_seq_len]
context_mask = context_mask.view(batch_size, 1, 1, -1)
if num_gen_seq > 1:
assert (
batch_size == 1
), f"num_gen_seq > 1 is only supported for a single prompt, got {len(prompt_tokens)} prompts"
log.debug(f"Generating {num_gen_seq} sequences with the same prompt")
assert (
num_gen_seq <= params.max_batch_size
), f"num_gen_seq={num_gen_seq} exceeds max_batch_size={params.max_batch_size}"
# repeat the prompt tokens for num_gen_seq times
prompt_tokens = prompt_tokens.repeat(num_gen_seq, 1)
assert prompt_tokens.shape == (
num_gen_seq,
prompt_len,
), f"prompt_tokens must be of shape (num_gen_seq, seq_len), got {prompt_tokens.shape}"
batch_size = len(prompt_tokens)
# create an empty tensor of the expected final shape and fill in the current tokens
empty = torch.empty(batch_size, total_len, dtype=prompt_tokens.dtype, device=prompt_tokens.device)
empty[:, :prompt_len] = prompt_tokens
seq = empty
input_pos = torch.arange(0, prompt_len, device="cuda")
if verbose:
prefill_start = time.time()
if images is not None:
images = images.to(device=prompt_tokens.device, dtype=torch.bfloat16)
prompt_token_embeddings = self.embed_vision_language_features(prompt_tokens, images)
else:
prompt_token_embeddings = None
if context is not None:
context = context.to(device=prompt_tokens.device, dtype=self.precision)
# Prefill stage
next_token = self.prefill(
self.model,
input_pos=input_pos,
tokens=prompt_tokens if prompt_token_embeddings is None else None,
token_embeddings=prompt_token_embeddings,
temperature=temperature,
top_k=top_k,
top_p=top_p,
context=context,
context_mask=context_mask,
)
if verbose:
prefill_time = time.time() - prefill_start
seq[:, [prompt_len]] = next_token.to(dtype=seq.dtype)
input_pos = torch.tensor([prompt_len], dtype=torch.long, device="cuda")
stop_tokens = self.tokenizer.stop_tokens if stop_tokens is None else stop_tokens
stop_tokens = torch.tensor(list(stop_tokens), dtype=torch.long, device="cuda")
if verbose:
decode_start = time.time()
# Decode stage
generated_tokens = decode_n_tokens(
self.model,
next_token.view(batch_size, -1),
input_pos,
max_gen_len - 1,
temperature=temperature,
top_k=top_k,
top_p=top_p,
stop_tokens=stop_tokens,
decode_one_token_function=self.decode_one_token,
context=context,
context_mask=context_mask,
)
gen_len = len(generated_tokens)
if verbose:
decode_time = time.time() - decode_start
prefill_throughput = prompt_len / prefill_time
decode_throughput = gen_len / decode_time
log.debug(f"[Prefill] Time: {prefill_time:.2f}s; Throughput: {prefill_throughput:.2f} tokens/s")
log.debug(f"[Decode] Time: {decode_time:.2f}s; Throughput: {decode_throughput:.2f} tokens/s")
generated_tokens = torch.cat(generated_tokens, dim=1)
log.debug(f"generated_tokens: {generated_tokens.shape}")
seq = seq[:, : prompt_len + 1 + gen_len]
seq[:, prompt_len + 1 :] = generated_tokens
if not echo:
seq = seq[:, prompt_len:]
torch.set_default_dtype(orig_precision) # Reset the default dtype to the original value
return seq, None
def embed_vision_language_features(self, input_ids: torch.Tensor, images: torch.tensor) -> torch.Tensor:
"""
Embed vision and language features into a combined representation.
Args:
input_ids (torch.Tensor): Input token IDs.
images (torch.tensor): Input images.
Returns:
torch.Tensor: Combined vision-language features.
Raises:
AssertionError: If vision encoder or mm projector is not initialized,
or if dimensions mismatch.
"""
# Ensure vision encoder and mm projector are initialized
assert self.vision_encoder is not None
assert self.mm_projector is not None
# Get image token ID and validate it
image_token_id = self.vision_encoder.image_token_id
assert isinstance(image_token_id, int) and image_token_id >= 0, f"Invalid image_token_id: {image_token_id}"
# Identify text and image locations in the input
text_locations = input_ids != image_token_id
image_locations = input_ids == image_token_id
# Process text features
text_features = self.model.tok_embeddings(input_ids[text_locations])
# Process image features
images = images.to(device=text_features.device, dtype=text_features.dtype)
vit_outputs = self.vision_encoder(images)
image_features = self.mm_projector(vit_outputs)
# Get dimensions
B, seq_len = input_ids.shape
N_total = B * seq_len
N_txt, D_txt = text_features.shape
N_img, N_patch, D_img = image_features.shape
# Reshape image features
image_features = image_features.reshape(N_img * N_patch, D_img)
# Validate dimensions
assert D_txt == D_img, f"Text features dim {D_txt} should be equal to image features dim {D_img}"
assert (
N_total == N_txt + N_img * N_patch
), f"seq_len {seq_len} should be equal to N_txt + N_img*N_Patch {(N_txt, N_img * N_patch, image_locations.sum().item())}"
# Combine text and image features
combined_features = torch.empty(
(B, seq_len, D_txt),
dtype=text_features.dtype,
device=text_features.device,
)
combined_features[text_locations, :] = text_features
combined_features[image_locations, :] = image_features
return combined_features
def state_dict(self, *args, **kwargs):
"""
Process the state dict (e.g., remove "_extra_state" keys imposed by TransformerEngine for FP8).
"""
state_dict = super().state_dict(*args, **kwargs)
return process_state_dict(state_dict)
def load_state_dict(self, state_dict: Dict[str, Any], strict: bool = True, assign: bool = False):
"""
Ignore the missing keys with substrings matching `substring_to_ignore` (e.g., "_extra_state" keys imposed by
TransformerEngine for FP8).
"""
state_dict = process_state_dict(state_dict)
missing_keys, unexpected_keys = super().load_state_dict(state_dict, strict=False, assign=assign)
actual_missing_keys = []
for key in missing_keys:
if not any(substring in key for substring in substrings_to_ignore):
actual_missing_keys.append(key)
if strict:
if len(actual_missing_keys) > 0 or len(unexpected_keys) > 0:
raise ValueError(f"Missing keys: {actual_missing_keys}\n\nUnexpected keys: {unexpected_keys}")
return _IncompatibleKeys(actual_missing_keys, unexpected_keys)