runtime_utils

Utilities for runtime benchmarking and model saving in Puzzletron.

This module provides classes and utility functions used for empirical runtime estimation of Transformer subblocks and for saving models and tokenizers in formats suitable for benchmarking with vLLM.

Classes

RuntimeConfig

Configuration for a vLLM latency benchmark run.

Functions

convert_config_to_vllm_anymodel	Convert a model to vLLM AnyModel format.
save_model	Save model weights as AnyModel and copy the tokenizer to output_path.
save_model_as_anymodel	Save a temporary vLLM-compatible AnyModel benchmark checkpoint.

class RuntimeConfig

Bases: object

Configuration for a vLLM latency benchmark run.

__init__(vocab_size, hidden_size, num_attention_heads, num_key_value_heads, descriptor, model_config_fields, tokenizer_path, repeat_block_n_times, prefill_seq_len, generation_seq_len, batch_size, num_iters, num_warmup_iters, extra_vllm_args=(), max_num_seqs=None, topology=RuntimeTopology(tensor_parallel_size=1, pipeline_parallel_size=1, data_parallel_size=1, prefill_context_parallel_size=1, decode_context_parallel_size=1, enable_expert_parallel=False, distributed_executor_backend='mp', gpu_group_size=1), estimator_schema='candidate_slope_v1', estimator_mode='homogeneous', effective_repeat_count=None, scaffold_policy='none', vllm_env=())

Parameters:

• vocab_size (int)

• hidden_size (int)

• num_attention_heads (int)

• num_key_value_heads (int)

• descriptor (type)

• model_config_fields (tuple[tuple[str, Any], ...])

• tokenizer_path (str)

• repeat_block_n_times (int)

• prefill_seq_len (int)

• generation_seq_len (int)

• batch_size (int)

• num_iters (int)

• num_warmup_iters (int)

• extra_vllm_args (tuple[str, ...])

• max_num_seqs (int | None)

• topology (RuntimeTopology)

• estimator_schema (str)

• estimator_mode (str)

• effective_repeat_count (int | None)

• scaffold_policy (str)

• vllm_env (tuple[tuple[str, str], ...])

Return type:

None

batch_size: int

descriptor: type

effective_repeat_count: int | None = None

estimator_mode: str = 'homogeneous'

estimator_schema: str = 'candidate_slope_v1'

extra_vllm_args: tuple[str, ...] = ()

generation_seq_len: int

hidden_size: int

max_num_seqs: int | None = None

model_config_fields: tuple[tuple[str, Any], ...]

model_config_value(key, default=None)

Return a descriptor-specific benchmark config value.

Parameters:

• key (str)

• default (Any)

Return type:

Any

num_attention_heads: int

num_iters: int

num_key_value_heads: int

num_warmup_iters: int

prefill_seq_len: int

repeat_block_n_times: int

scaffold_policy: str = 'none'

tokenizer_path: str

topology: RuntimeTopology = RuntimeTopology(tensor_parallel_size=1, pipeline_parallel_size=1, data_parallel_size=1, prefill_context_parallel_size=1, decode_context_parallel_size=1, enable_expert_parallel=False, distributed_executor_backend='mp', gpu_group_size=1)

vllm_env: tuple[tuple[str, str], ...] = ()

vocab_size: int

convert_config_to_vllm_anymodel(config_dir)

Convert a model to vLLM AnyModel format.

Parameters:

config_dir (Path)

save_model(model, tokenizer_path, output_path, descriptor)

Save model weights as AnyModel and copy the tokenizer to output_path.

Parameters:

• model (PreTrainedModel)

• tokenizer_path (Path)

• output_path (Path)

• descriptor (type)

Return type:

None

save_model_as_anymodel(model, output_dir, descriptor, runtime_descriptor=None)

Save a temporary vLLM-compatible AnyModel benchmark checkpoint.

Parameters:

output_dir (Path)