diff --git a/inference/generate.py b/inference/generate.py
index fbf3ab8..a76b66b 100644
--- a/inference/generate.py
+++ b/inference/generate.py
@@ -1,31 +1,91 @@
import os
import json
+import logging
from argparse import ArgumentParser
-from typing import List
+from pathlib import Path
+from typing import List, Optional, Dict, Tuple
+from contextlib import nullcontext
import torch
import torch.distributed as dist
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, AutoConfig
from safetensors.torch import load_model
+from tqdm import tqdm
from model import Transformer, ModelArgs
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
-def sample(logits, temperature: float = 1.0):
+# Constants
+DEFAULT_EOS_TOKEN = "</s>"
+MAX_SEQ_LEN_WARNING_THRESHOLD = 0.9
+TORCH_DTYPE = torch.bfloat16
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+def setup_distributed() -> Tuple[int, int, int]:
+ """Initialize distributed training environment."""
+ world_size = int(os.getenv("WORLD_SIZE", "1"))
+ rank = int(os.getenv("RANK", "0"))
+ local_rank = int(os.getenv("LOCAL_RANK", "0"))
+
+ if world_size > 1:
+ dist.init_process_group(
+ backend="nccl" if torch.cuda.is_available() else "gloo",
+ timeout=timedelta(minutes=5)
+ )
+ logger.info(f"Initialized process group (rank {rank}/{world_size})")
+
+ torch.cuda.set_device(local_rank)
+ return world_size, rank, local_rank
+
+def validate_paths(ckpt_path: Path, config_path: Path) -> None:
+ """Validate model checkpoint and config paths."""
+ if not ckpt_path.exists():
+ raise FileNotFoundError(f"Checkpoint directory {ckpt_path} not found")
+ if not config_path.exists():
+ raise FileNotFoundError(f"Config file {config_path} not found")
+
+def load_model_config(config_path: Path) -> ModelArgs:
+ """Load and validate model configuration."""
+ try:
+ with open(config_path) as f:
+ config_data = json.load(f)
+ return ModelArgs(**config_data)
+ except (json.JSONDecodeError, TypeError) as e:
+ logger.error(f"Invalid model config: {str(e)}")
+ raise
+
+def initialize_model(args: ModelArgs, device: str) -> Transformer:
+ """Initialize model with proper device placement and dtype."""
+ with torch.device(device):
+ model = Transformer(args)
+ model.to(TORCH_DTYPE)
+ model.eval()
+ return model
+
+def sample(logits: torch.Tensor, temperature: float = 1.0, top_k: int = 50) -> torch.Tensor:
"""
- Samples a token from the logits using temperature scaling.
-
+ Sample token from logits with temperature and top-k filtering.
+
Args:
- logits (torch.Tensor): The logits tensor for token predictions.
- temperature (float, optional): Temperature for scaling logits. Defaults to 1.0.
-
+ logits: Unnormalized log probabilities (batch_size, vocab_size)
+ temperature: Sampling temperature (0.0 = greedy)
+ top_k: Top-k tokens to consider (0 = no filtering)
+
Returns:
- torch.Tensor: The sampled token.
+ Sampled token indices (batch_size, 1)
"""
- logits = logits / max(temperature, 1e-5)
- probs = torch.softmax(logits, dim=-1)
- return probs.div_(torch.empty_like(probs).exponential_(1)).argmax(dim=-1)
-
+ if temperature <= 0:
+ return logits.argmax(dim=-1)
+
+ if top_k > 0:
+ v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+ logits[logits < v[:, [-1]]] = -float('inf')
+
+ probs = torch.softmax(logits / max(temperature, 1e-5), dim=-1)
+ return torch.multinomial(probs, num_samples=1).squeeze(1)
@torch.inference_mode()
def generate(
@@ -33,153 +93,260 @@ def generate(
prompt_tokens: List[List[int]],
max_new_tokens: int,
eos_id: int,
- temperature: float = 1.0
+ temperature: float = 1.0,
+ top_k: int = 50,
+ repetition_penalty: float = 1.1
) -> List[List[int]]:
"""
- Generates new tokens based on the given prompt tokens using the specified model.
-
+ Generate text with dynamic sequence length management.
+
Args:
- model (Transformer): The transformer model used for token generation.
- prompt_tokens (List[List[int]]): A list of lists containing the prompt tokens for each sequence.
- max_new_tokens (int): The maximum number of new tokens to generate.
- eos_id (int): The end-of-sequence token ID.
- temperature (float, optional): The temperature value for sampling. Defaults to 1.0.
-
+ model: Initialized transformer model
+ prompt_tokens: List of tokenized prompts
+ max_new_tokens: Maximum new tokens to generate
+ eos_id: End-of-sequence token ID
+ temperature: Sampling temperature
+ top_k: Top-k sampling parameter
+ repetition_penalty: Penalty for repeated tokens
+
Returns:
- List[List[int]]: A list of lists containing the generated tokens for each sequence.
+ List of generated token sequences
"""
- prompt_lens = [len(t) for t in prompt_tokens]
- assert max(prompt_lens) <= model.max_seq_len
- total_len = min(model.max_seq_len, max_new_tokens + max(prompt_lens))
- tokens = torch.full((len(prompt_tokens), total_len), -1, dtype=torch.long, device="cuda")
- for i, t in enumerate(prompt_tokens):
- tokens[i, :len(t)] = torch.tensor(t, dtype=torch.long, device="cuda")
+ # Initialize generation state
+ batch_size = len(prompt_tokens)
+ device = next(model.parameters()).device
+ max_seq_len = model.max_seq_len
+ prompt_lens = [len(p) for p in prompt_tokens]
+
+ # Validate input lengths
+ if max(prompt_lens) + max_new_tokens > max_seq_len:
+ logger.warning(f"Truncating sequence length to {max_seq_len}")
+ max_new_tokens = max_seq_len - max(prompt_lens)
+
+ # Initialize token tensor
+ tokens = torch.full((batch_size, max_seq_len), -1, dtype=torch.long, device=device)
+ for i, seq in enumerate(prompt_tokens):
+ tokens[i, :len(seq)] = torch.tensor(seq, device=device)
+
+ # Generation loop
prev_pos = 0
- finished = torch.tensor([False] * len(prompt_tokens), device="cuda")
+ finished = torch.zeros(batch_size, dtype=torch.bool, device=device)
prompt_mask = tokens != -1
- for cur_pos in range(min(prompt_lens), total_len):
- logits = model.forward(tokens[:, prev_pos:cur_pos], prev_pos)
- if temperature > 0:
- next_token = sample(logits, temperature)
- else:
- next_token = logits.argmax(dim=-1)
- next_token = torch.where(prompt_mask[:, cur_pos], tokens[:, cur_pos], next_token)
- tokens[:, cur_pos] = next_token
- finished |= torch.logical_and(~prompt_mask[:, cur_pos], next_token == eos_id)
- prev_pos = cur_pos
- if finished.all():
- break
- completion_tokens = []
- for i, toks in enumerate(tokens.tolist()):
- toks = toks[prompt_lens[i]:prompt_lens[i]+max_new_tokens]
- if eos_id in toks:
- toks = toks[:toks.index(eos_id)]
- completion_tokens.append(toks)
- return completion_tokens
+ progress_bar = tqdm(total=max_new_tokens, desc="Generating", disable=not logger.isEnabledFor(logging.INFO))
+
+ try:
+ for cur_pos in range(max(prompt_lens), min(max_seq_len, max(prompt_lens) + max_new_tokens)):
+ # Model forward pass
+ logits = model(tokens[:, prev_pos:cur_pos], prev_pos)
+
+ # Apply repetition penalty
+ if repetition_penalty != 1.0:
+ for idx in range(batch_size):
+ unique_tokens, counts = torch.unique(tokens[idx], return_counts=True)
+ logits[idx, unique_tokens] /= counts.float() ** (repetition_penalty - 1.0)
+
+ # Sample next tokens
+ next_tokens = sample(logits[:, -1], temperature, top_k)
+
+ # Update tokens
+ tokens[:, cur_pos] = torch.where(
+ prompt_mask[:, cur_pos],
+ tokens[:, cur_pos],
+ next_tokens
+ )
+
+ # Update completion status
+ finished |= (~prompt_mask[:, cur_pos] & (next_tokens == eos_id))
+ prev_pos = cur_pos
+ progress_bar.update(1)
+
+ if finished.all():
+ break
+ finally:
+ progress_bar.close()
+
+ # Process outputs
+ return [seq[pl:pl+max_new_tokens].tolist() for pl, seq in zip(prompt_lens, tokens)]
+def interactive_loop(
+ model: Transformer,
+ tokenizer: AutoTokenizer,
+ world_size: int,
+ rank: int,
+ max_new_tokens: int,
+ temperature: float
+) -> None:
+ """Interactive chat interface with history management."""
+ messages = []
+ eos_id = tokenizer.eos_token_id or tokenizer.convert_tokens_to_ids(DEFAULT_EOS_TOKEN)
+
+ while True:
+ try:
+ # Distributed input handling
+ if world_size > 1:
+ if rank == 0:
+ prompt = input("\nUser: ")
+ dist.broadcast_object_list([prompt], src=0)
+ else:
+ prompt = None
+ dist.broadcast_object_list([prompt], src=0)
+
+ if prompt == "/exit":
+ break
+ else:
+ prompt = input("\nUser: ")
+
+ # Command handling
+ if prompt == "/exit":
+ break
+ if prompt == "/clear":
+ messages.clear()
+ logger.info("History cleared")
+ continue
+
+ # Tokenize and generate
+ messages.append({"role": "user", "content": prompt})
+ prompt_tokens = tokenizer.apply_chat_template(
+ messages,
+ add_generation_prompt=True,
+ truncation=True,
+ max_length=model.max_seq_len - max_new_tokens
+ )
+
+ completion_tokens = generate(
+ model,
+ [prompt_tokens],
+ max_new_tokens,
+ eos_id,
+ temperature
+ )[0]
+
+ # Decode and update history
+ completion = tokenizer.decode(completion_tokens, skip_special_tokens=True)
+ messages.append({"role": "assistant", "content": completion})
+ print(f"\nAssistant: {completion}")
+
+ except KeyboardInterrupt:
+ logger.info("\nExiting...")
+ break
+ except Exception as e:
+ logger.error(f"Generation error: {str(e)}")
+ messages.pop() # Remove failed prompt
+
+def batch_process(
+ model: Transformer,
+ tokenizer: AutoTokenizer,
+ input_file: Path,
+ max_new_tokens: int,
+ temperature: float
+) -> None:
+ """Batch processing mode with progress tracking."""
+ try:
+ with open(input_file) as f:
+ prompts = [line.strip() for line in f if line.strip()]
+
+ if not prompts:
+ raise ValueError("Input file is empty")
+
+ # Tokenize with parallel processing
+ tokenizer_fn = lambda p: tokenizer.apply_chat_template(
+ [{"role": "user", "content": p}],
+ add_generation_prompt=True,
+ truncation=True,
+ max_length=model.max_seq_len - max_new_tokens
+ )
+ prompt_tokens = [tokenizer_fn(p) for p in tqdm(prompts, desc="Tokenizing")]
+
+ # Generate in batches
+ completions = []
+ for i in tqdm(range(0, len(prompt_tokens), model.args.max_batch_size)):
+ batch = prompt_tokens[i:i+model.args.max_batch_size]
+ completions += generate(model, batch, max_new_tokens, tokenizer.eos_token_id, temperature)
+
+ # Decode and print
+ for prompt, tokens in zip(prompts, completions):
+ completion = tokenizer.decode(tokens, skip_special_tokens=True)
+ print(f"\nPrompt: {prompt}\nCompletion: {completion}\n{'='*50}")
+
+ except Exception as e:
+ logger.error(f"Batch processing failed: {str(e)}")
+ raise
def main(
ckpt_path: str,
- config: str,
+ config_path: str,
input_file: str = "",
interactive: bool = True,
- max_new_tokens: int = 100,
- temperature: float = 1.0,
+ max_new_tokens: int = 200,
+ temperature: float = 0.2
) -> None:
- """
- Main function to load the model and perform interactive or batch text generation.
-
- Args:
- ckpt_path (str): Path to the model checkpoint directory.
- config (str): Path to the model configuration file.
- input_file (str, optional): Path to a file containing input prompts. Defaults to "".
- interactive (bool, optional): Whether to run in interactive mode. Defaults to True.
- max_new_tokens (int, optional): Maximum number of new tokens to generate. Defaults to 100.
- temperature (float, optional): Temperature for sampling. Defaults to 1.0.
- """
- world_size = int(os.getenv("WORLD_SIZE", "1"))
- rank = int(os.getenv("RANK", "0"))
- local_rank = int(os.getenv("LOCAL_RANK", "0"))
- if world_size > 1:
- dist.init_process_group("nccl")
- global print
- if rank != 0:
- print = lambda *_, **__: None
- torch.cuda.set_device(local_rank)
- torch.set_default_dtype(torch.bfloat16)
- torch.set_num_threads(8)
- torch.manual_seed(965)
- with open(config) as f:
- args = ModelArgs(**json.load(f))
- print(args)
- with torch.device("cuda"):
- model = Transformer(args)
- tokenizer = AutoTokenizer.from_pretrained(ckpt_path)
- tokenizer.decode(generate(model, [tokenizer.encode("DeepSeek")], 2, -1, 1.)[0])
- load_model(model, os.path.join(ckpt_path, f"model{rank}-mp{world_size}.safetensors"))
-
- if interactive:
- messages = []
- while True:
- if world_size == 1:
- prompt = input(">>> ")
- elif rank == 0:
- prompt = input(">>> ")
- objects = [prompt]
- dist.broadcast_object_list(objects, 0)
- else:
- objects = [None]
- dist.broadcast_object_list(objects, 0)
- prompt = objects[0]
- if prompt == "/exit":
- break
- elif prompt == "/clear":
- messages.clear()
- continue
- messages.append({"role": "user", "content": prompt})
- prompt_tokens = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
- completion_tokens = generate(model, [prompt_tokens], max_new_tokens, tokenizer.eos_token_id, temperature)
- completion = tokenizer.decode(completion_tokens[0], skip_special_tokens=True)
- print(completion)
- messages.append({"role": "assistant", "content": completion})
- else:
- with open(input_file) as f:
- prompts = [line.strip() for line in f.readlines()]
- assert len(prompts) <= args.max_batch_size
- prompt_tokens = [tokenizer.apply_chat_template([{"role": "user", "content": prompt}], add_generation_prompt=True) for prompt in prompts]
- completion_tokens = generate(model, prompt_tokens, max_new_tokens, tokenizer.eos_token_id, temperature)
- completions = tokenizer.batch_decode(completion_tokens, skip_special_tokens=True)
- for prompt, completion in zip(prompts, completions):
- print("Prompt:", prompt)
- print("Completion:", completion)
- print()
-
- if world_size > 1:
- dist.destroy_process_group()
-
+ """Main execution flow with proper resource management."""
+ # Distributed setup
+ world_size, rank, local_rank = setup_distributed()
+
+ try:
+ # Path validation
+ ckpt_dir = Path(ckpt_path)
+ config_file = Path(config_path)
+ validate_paths(ckpt_dir, config_file)
+
+ # Model initialization
+ model_args = load_model_config(config_file)
+ model = initialize_model(model_args, DEVICE)
+ load_model(model, ckpt_dir / f"model{rank}-mp{world_size}.safetensors")
+
+ # Tokenizer setup
+ tokenizer = AutoTokenizer.from_pretrained(
+ ckpt_dir,
+ use_fast=True,
+ trust_remote_code=True
+ )
+
+ # Generation mode selection
+ if interactive:
+ interactive_loop(model, tokenizer, world_size, rank, max_new_tokens, temperature)
+ else:
+ batch_process(model, tokenizer, Path(input_file), max_new_tokens, temperature)
+
+ finally:
+ if world_size > 1:
+ dist.destroy_process_group()
if __name__ == "__main__":
- """
- Command-line interface for distributed text generation.
-
- Arguments:
- --ckpt-path (str): Path to the model checkpoint directory.
- --config (str): Path to the model configuration file.
- --input-file (str, optional): File containing prompts for batch processing.
- --interactive (bool, optional): Enable interactive mode for generating text.
- --max-new-tokens (int, optional): Maximum number of new tokens to generate. Defaults to 200.
- --temperature (float, optional): Temperature for sampling. Defaults to 0.2.
-
- Raises:
- AssertionError: If neither input-file nor interactive mode is specified.
- """
- parser = ArgumentParser()
- parser.add_argument("--ckpt-path", type=str, required=True)
- parser.add_argument("--config", type=str, required=True)
- parser.add_argument("--input-file", type=str, default="")
- parser.add_argument("--interactive", action="store_true")
- parser.add_argument("--max-new-tokens", type=int, default=200)
- parser.add_argument("--temperature", type=float, default=0.2)
+ parser = ArgumentParser(description="Distributed Transformer Text Generation")
+ parser.add_argument("--ckpt-path", type=str, required=True,
+ help="Path to model checkpoint directory")
+ parser.add_argument("--config", type=str, required=True,
+ help="Path to model config JSON file")
+ parser.add_argument("--input-file", type=str, default="",
+ help="Path to input file for batch processing")
+ parser.add_argument("--interactive", action="store_true",
+ help="Enable interactive chat mode")
+ parser.add_argument("--max-new-tokens", type=int, default=200,
+ help="Maximum new tokens to generate")
+ parser.add_argument("--temperature", type=float, default=0.2,
+ help="Sampling temperature (0.0 = greedy)")
+ parser.add_argument("--log-level", choices=["DEBUG", "INFO", "WARNING"], default="INFO",
+ help="Set logging verbosity")
+
args = parser.parse_args()
- assert args.input_file or args.interactive
- main(args.ckpt_path, args.config, args.input_file, args.interactive, args.max_new_tokens, args.temperature)
+
+ # Validate arguments
+ if not args.interactive and not args.input_file:
+ parser.error("Must specify either --interactive or --input-file")
+
+ # Configure logging
+ logger.setLevel(args.log_level)
+
+ try:
+ main(
+ args.ckpt_path,
+ args.config,
+ args.input_file,
+ args.interactive,
+ args.max_new_tokens,
+ args.temperature
+ )
+ except Exception as e:
+ logger.critical(f"Critical error: {str(e)}", exc_info=True)
+ exit(1)