diff --git a/assets/docs/architecture.md b/assets/docs/architecture.md index ced6f88..0be1bfc 100644 --- a/assets/docs/architecture.md +++ b/assets/docs/architecture.md @@ -28,17 +28,17 @@ classDiagram class AutoRegressiveLMConfig { +Optional[int] vocab_size - +Optional[int] dim - +Optional[int] n_layers - +Optional[float] norm_eps - +Optional[int] dim_ffn - +Optional[bool] tie_weight + +Optional[int] hidden_size + +Optional[int] num_hidden_layers + +Optional[float] rms_norm_eps + +Optional[int] intermediate_size + +Optional[bool] tie_word_embeddings +Optional[dict] rope_scaling - +Optional[int] max_len + +Optional[int] max_position_embeddings +Optional[float] rope_theta +str attn_type - +Optional[int] n_heads - +Optional[int] n_kv_heads + +Optional[int] num_attention_heads + +Optional[int] num_key_value_heads +Optional[bool] use_qk_norm +Optional[bool] use_gated_attention +Optional[int] kv_lora_rank @@ -53,15 +53,15 @@ classDiagram class EncoderConfig { +Optional[int] vocab_size - +Optional[int] dim - +Optional[int] n_layers - +Optional[float] norm_eps - +Optional[int] dim_ffn - +Optional[int] max_len + +Optional[int] hidden_size + +Optional[int] num_hidden_layers + +Optional[float] rms_norm_eps + +Optional[int] intermediate_size + +Optional[int] max_position_embeddings +Optional[float] rope_theta +str attn_type - +Optional[int] n_heads - +Optional[int] n_kv_heads + +Optional[int] num_attention_heads + +Optional[int] num_key_value_heads +Optional[bool] use_qk_norm +str ffn_type +Optional[dict] rope_scaling diff --git a/assets/docs/dataflow.md b/assets/docs/dataflow.md index fa5bca4..0c676e6 100644 --- a/assets/docs/dataflow.md +++ b/assets/docs/dataflow.md @@ -85,7 +85,7 @@ All backends normalise tensors into `Store._data[Dict[str, List[Tensor]]]` + `St ``` DatasetFactory.load(train_type, load_path, window_size, stride=None, storage_type=None, tokenizer_path=None, - max_len=2048, store=None) + max_position_embeddings=2048, store=None) → BaseDataset.load(load_path, storage_type=None) → detect_format(load_path) → StoreFactory.create(storage_type) diff --git a/assets/docs/inference.md b/assets/docs/inference.md index 8566b4f..8eb2a0c 100644 --- a/assets/docs/inference.md +++ b/assets/docs/inference.md @@ -32,7 +32,7 @@ ContiguousCache (simple contiguous per-slot cache) ├── ContiguousCacheView bundles k/v tensors + slot indices for attention layers ``` -Created by default when no cache is passed to `InferenceScheduler`. Each task occupies a fixed slot of `[max_seq_len, n_kv_heads, head_dim]`. Simple and efficient for small-to-medium batch sizes. +Created by default when no cache is passed to `InferenceScheduler`. Each task occupies a fixed slot of `[max_seq_len, num_key_value_heads, head_dim]`. Simple and efficient for small-to-medium batch sizes. ### PageCache (paged with prefix sharing) @@ -42,7 +42,7 @@ PageCache (paged KV cache with prefix sharing, alternative) │ ├── Allocator bitmask-based page allocator + ref-count + LRU │ └── PrefixCache hash-based prefix matching (page_hash via polynomial hash) ├── TaskTable maps task_id → page_table + cached token count - ├── Storage k_cache / v_cache tensors (n_layers × n_pages × page_size × n_kv_heads × head_dim) + ├── Storage k_cache / v_cache tensors (num_hidden_layers × n_pages × page_size × num_key_value_heads × head_dim) └── PageCacheView bundles Storage + page_table + total_len for attention layers ``` diff --git a/assets/docs/params.md b/assets/docs/params.md index 65b75ef..cfb6ed9 100644 --- a/assets/docs/params.md +++ b/assets/docs/params.md @@ -44,7 +44,7 @@ Combined optimizer: matrix parameters via **Muon**, non-matrix via **AdamW** (`f | Parameter | Description | Default | |-----------|-------------|---------| -| `--window_size` | Max input sequence length | model config `max_len` | +| `--window_size` | Max input sequence length | model config `max_position_embeddings` | | `--stride` | Stride for sliding window over sequences | None | | `--random_seed` | Random seed for reproducibility | 3407 | | `--num_workers` | DataLoader worker processes | 4 | @@ -186,7 +186,7 @@ See [Inference Guide](inference.md) for HTTP API documentation. | `--top_k` | int | `30` | Top-k filtering | | `--top_p` | float | `0.95` | Nucleus sampling threshold | | `--batch_size` | int | `1` | Batch size for generation | -| `--max_tokens` | int | model config `max_len` | Maximum tokens to generate | +| `--max_tokens` | int | model config `max_position_embeddings` | Maximum tokens to generate | Usage: ```bash diff --git a/astrai/config/model_config.py b/astrai/config/model_config.py index d9b8e5c..a74842a 100644 --- a/astrai/config/model_config.py +++ b/astrai/config/model_config.py @@ -29,19 +29,19 @@ class AutoRegressiveLMConfig(BaseModelConfig): """Configuration for autoregressive language model.""" vocab_size: Optional[int] = None - dim: Optional[int] = None - n_layers: Optional[int] = None - norm_eps: Optional[float] = None - dim_ffn: Optional[int] = None - tie_weight: Optional[bool] = None + hidden_size: Optional[int] = None + num_hidden_layers: Optional[int] = None + rms_norm_eps: Optional[float] = None + intermediate_size: Optional[int] = None + tie_word_embeddings: Optional[bool] = None - max_len: Optional[int] = None + max_position_embeddings: Optional[int] = None rope_theta: Optional[float] = None rope_scaling: Optional[dict] = None attn_type: str = "gqa" - n_heads: Optional[int] = None - n_kv_heads: Optional[int] = None + num_attention_heads: Optional[int] = None + num_key_value_heads: Optional[int] = None use_qk_norm: Optional[bool] = None use_gated_attention: Optional[bool] = None @@ -62,18 +62,18 @@ class EncoderConfig(BaseModelConfig): """Configuration for embedding encoder model.""" vocab_size: Optional[int] = None - dim: Optional[int] = None - n_layers: Optional[int] = None - norm_eps: Optional[float] = None - dim_ffn: Optional[int] = None + hidden_size: Optional[int] = None + num_hidden_layers: Optional[int] = None + rms_norm_eps: Optional[float] = None + intermediate_size: Optional[int] = None - max_len: Optional[int] = None + max_position_embeddings: Optional[int] = None rope_theta: Optional[float] = None rope_scaling: Optional[dict] = None attn_type: str = "gqa" - n_heads: Optional[int] = None - n_kv_heads: Optional[int] = None + num_attention_heads: Optional[int] = None + num_key_value_heads: Optional[int] = None use_qk_norm: Optional[bool] = None use_gated_attention: Optional[bool] = None diff --git a/astrai/inference/core/scheduler.py b/astrai/inference/core/scheduler.py index 92cf89f..2b1ccab 100644 --- a/astrai/inference/core/scheduler.py +++ b/astrai/inference/core/scheduler.py @@ -32,26 +32,26 @@ class InferenceScheduler: if max_seq_len is not None: self.max_seq_len = max_seq_len - elif config.max_len is not None: - self.max_seq_len = config.max_len + elif config.max_position_embeddings is not None: + self.max_seq_len = config.max_position_embeddings else: raise ValueError( "max_seq_len must be provided either as argument " - "or in model config (config.max_len)" + "or in model config (config.max_position_embeddings)" ) self.device = device or next(model.parameters()).device self.dtype = dtype or next(model.parameters()).dtype - head_dim = config.dim // config.n_heads + head_dim = config.hidden_size // config.num_attention_heads if cache is not None: self._cache = cache else: self._cache = ContiguousCache( - config.n_layers, + config.num_hidden_layers, max_batch_size, self.max_seq_len, - config.n_kv_heads, + config.num_key_value_heads, head_dim, self.device, self.dtype, diff --git a/astrai/model/components/decoder_block.py b/astrai/model/components/decoder_block.py index e686fd0..7abc96f 100644 --- a/astrai/model/components/decoder_block.py +++ b/astrai/model/components/decoder_block.py @@ -14,10 +14,18 @@ class DecoderBlock(nn.Module): def __init__(self, config, layer_id: int): super().__init__() cfg = asdict(config) - cfg["down_init_std"] = 0.02 / (2 * config.n_layers) ** 0.5 + cfg.update( + dim=config.hidden_size, + dim_ffn=config.intermediate_size, + n_layers=config.num_hidden_layers, + n_heads=config.num_attention_heads, + n_kv_heads=config.num_key_value_heads, + norm_eps=config.rms_norm_eps, + down_init_std=0.02 / (2 * config.num_hidden_layers) ** 0.5, + ) self.attention = AttnFactory.create(config.attn_type, **cfg, layer_id=layer_id) - self.input_norm = RMSNorm(config.dim, config.norm_eps) - self.post_attention_norm = RMSNorm(config.dim, config.norm_eps) + self.input_norm = RMSNorm(config.hidden_size, config.rms_norm_eps) + self.post_attention_norm = RMSNorm(config.hidden_size, config.rms_norm_eps) self.mlp = FFNFactory.create(config.ffn_type, **cfg) def forward( diff --git a/astrai/model/encoder.py b/astrai/model/encoder.py index cd2b5c2..3d2bd4a 100644 --- a/astrai/model/encoder.py +++ b/astrai/model/encoder.py @@ -18,20 +18,28 @@ class EmbeddingEncoder(AutoModel): def __init__(self, config: EncoderConfig): super().__init__(config) self.config = config - rope_dim = config.dim // config.n_heads + rope_dim = config.hidden_size // config.num_attention_heads rope_base = config.rope_theta if config.rope_theta is not None else 10000 self.rotary_embedding = RotaryEmbedding( - rope_dim, config.max_len, rope_base, rope_scaling=config.rope_scaling + rope_dim, + config.max_position_embeddings, + rope_base, + rope_scaling=config.rope_scaling, ) self.embed_tokens = Embedding( - config.vocab_size, config.dim, neftune_alpha=config.neftune_alpha + config.vocab_size, + config.hidden_size, + neftune_alpha=config.neftune_alpha, ) self.layers = nn.ModuleList( - [DecoderBlock(config, layer_id) for layer_id in range(config.n_layers)] + [ + DecoderBlock(config, layer_id) + for layer_id in range(config.num_hidden_layers) + ] ) - self.norm = RMSNorm(config.dim, config.norm_eps) + self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps) self.pooling_type = config.pooling_type or "mean" self.normalize_embeddings = config.normalize_embeddings or False diff --git a/astrai/model/transformer.py b/astrai/model/transformer.py index 117d757..c194d52 100644 --- a/astrai/model/transformer.py +++ b/astrai/model/transformer.py @@ -36,24 +36,32 @@ class AutoRegressiveLM(AutoModel): rope_dim = ( config.qk_rope_head_dim if config.attn_type == "mla" - else config.dim // config.n_heads + else config.hidden_size // config.num_attention_heads ) rope_base = config.rope_theta if config.rope_theta is not None else 10000 self.rotary_embedding = RotaryEmbedding( - rope_dim, config.max_len, rope_base, rope_scaling=config.rope_scaling + rope_dim, + config.max_position_embeddings, + rope_base, + rope_scaling=config.rope_scaling, ) self.embed_tokens = Embedding( - config.vocab_size, config.dim, neftune_alpha=config.neftune_alpha + config.vocab_size, + config.hidden_size, + neftune_alpha=config.neftune_alpha, ) self.layers = nn.ModuleList( - [DecoderBlock(config, layer_id) for layer_id in range(config.n_layers)] + [ + DecoderBlock(config, layer_id) + for layer_id in range(config.num_hidden_layers) + ] ) - self.norm = RMSNorm(config.dim, config.norm_eps) - self.lm_head = Linear(config.dim, config.vocab_size) + self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps) + self.lm_head = Linear(config.hidden_size, config.vocab_size) - if self.config.tie_weight is True: + if self.config.tie_word_embeddings is True: self.lm_head.weight = self.embed_tokens.weight self.apply(self._init_weights) @@ -68,7 +76,7 @@ class AutoRegressiveLM(AutoModel): state_dict = dict(state_dict) - if self.config.tie_weight is True: + if self.config.tie_word_embeddings is True: # same tensor for embed and lm_head if embed_key in state_dict: state_dict[lm_head_key] = state_dict[embed_key] @@ -84,7 +92,7 @@ class AutoRegressiveLM(AutoModel): destination=destination, prefix=prefix, keep_vars=keep_vars ) - if self.config.tie_weight is True: + if self.config.tie_word_embeddings is True: lm_head_key = prefix + "lm_head.weight" if lm_head_key in state_dict: del state_dict[lm_head_key] diff --git a/astrai/trainer/train_context.py b/astrai/trainer/train_context.py index d67c43d..c089846 100644 --- a/astrai/trainer/train_context.py +++ b/astrai/trainer/train_context.py @@ -240,7 +240,7 @@ class TrainContextBuilder: group_size = strategy_kwargs.get("group_size", 1) rollout_batch_size = group_size * max(1, cfg.batch_per_device) - max_seq_len = getattr(context.model.config, "max_len", None) + max_seq_len = getattr(context.model.config, "max_position_embeddings", None) scheduler = InferenceScheduler( model=context.model, diff --git a/scripts/eval/evaluate_mmlu.py b/scripts/eval/evaluate_mmlu.py index d2db818..9263c7d 100644 --- a/scripts/eval/evaluate_mmlu.py +++ b/scripts/eval/evaluate_mmlu.py @@ -185,7 +185,7 @@ def choice_logprob( choice_text = choice_letter choice_ids = tokenizer.encode(choice_text, add_special_tokens=False) input_ids = context_ids + choice_ids - max_len = model.config.max_len + max_len = model.config.max_position_embeddings if len(input_ids) > max_len: overflow = len(input_ids) - max_len input_ids = input_ids[overflow:] diff --git a/scripts/tools/benchmark.py b/scripts/tools/benchmark.py index 59930d1..92cad48 100644 --- a/scripts/tools/benchmark.py +++ b/scripts/tools/benchmark.py @@ -119,15 +119,15 @@ class GenerationBenchmark: dtype=torch.long, ) - head_dim = self.config.dim // self.config.n_heads + head_dim = self.config.hidden_size // self.config.num_attention_heads max_seq = prompt_length + gen_length if self.cache_type == "contiguous": cache = ContiguousCache( - self.config.n_layers, + self.config.num_hidden_layers, batch_size, max_seq, - self.config.n_kv_heads, + self.config.num_key_value_heads, head_dim, self.device, self.dtype, @@ -136,10 +136,10 @@ class GenerationBenchmark: page_size = 128 n_pages = (max_seq + page_size - 1) // page_size * batch_size cache = PageCache( - self.config.n_layers, + self.config.num_hidden_layers, n_pages, page_size, - self.config.n_kv_heads, + self.config.num_key_value_heads, head_dim, self.device, self.dtype, @@ -262,13 +262,13 @@ if __name__ == "__main__": config = AutoRegressiveLMConfig( vocab_size=10000, - dim=1536, - n_heads=24, - n_kv_heads=4, - dim_ffn=6912, - max_len=2048, - n_layers=24, - norm_eps=1e-5, + hidden_size=1536, + num_attention_heads=24, + num_key_value_heads=4, + intermediate_size=6912, + max_position_embeddings=2048, + num_hidden_layers=24, + rms_norm_eps=1e-5, ) benchmark = GenerationBenchmark( diff --git a/scripts/tools/generate.py b/scripts/tools/generate.py index 45eaaf3..ec8abbc 100644 --- a/scripts/tools/generate.py +++ b/scripts/tools/generate.py @@ -56,7 +56,7 @@ def processor( print(f" {len(prompts)} prompts loaded\n") if max_tokens is None: - max_tokens = model.config.max_len + max_tokens = model.config.max_position_embeddings chunk_size = max(1, batch_size) @@ -185,7 +185,10 @@ if __name__ == "__main__": "--max_tokens", type=int, default=None, - help="Maximum tokens to generate (default: model config max_len).", + help=( + "Maximum tokens to generate " + "(default: model config max_position_embeddings)." + ), ) parser.add_argument( "--cache_len", diff --git a/scripts/tools/train.py b/scripts/tools/train.py index 0a223e2..592df4e 100644 --- a/scripts/tools/train.py +++ b/scripts/tools/train.py @@ -480,7 +480,7 @@ def train( config.neftune_alpha = neftune_alpha if window_size is None: - window_size = config.max_len + window_size = config.max_position_embeddings strategy_kwargs = { "beta": kwargs.pop("dpo_beta"), diff --git a/tests/conftest.py b/tests/conftest.py index 087b87c..12ce6fe 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -107,13 +107,13 @@ def test_model(): """Session-scoped small AutoRegressiveLM model, created once.""" config = AutoRegressiveLMConfig( vocab_size=1000, - dim=8, - n_heads=2, - n_kv_heads=1, - dim_ffn=16, - max_len=64, - n_layers=2, - norm_eps=1e-5, + hidden_size=8, + num_attention_heads=2, + num_key_value_heads=1, + intermediate_size=16, + max_position_embeddings=64, + num_hidden_layers=2, + rms_norm_eps=1e-5, ) device = "cuda" if torch.cuda.is_available() else "cpu" model = AutoRegressiveLM(config).to(device=device) @@ -137,13 +137,13 @@ def base_test_env(test_model, test_tokenizer): json.dump( { "vocab_size": 1000, - "dim": 8, - "n_heads": 2, - "n_kv_heads": 1, - "dim_ffn": 16, - "max_len": 64, - "n_layers": 2, - "norm_eps": 1e-5, + "hidden_size": 8, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "intermediate_size": 16, + "max_position_embeddings": 64, + "num_hidden_layers": 2, + "rms_norm_eps": 1e-5, }, f, ) diff --git a/tests/inference/test_scheduler.py b/tests/inference/test_scheduler.py index 5125f69..91b8502 100644 --- a/tests/inference/test_scheduler.py +++ b/tests/inference/test_scheduler.py @@ -14,11 +14,11 @@ def mock_model_and_tokenizer(): """Create mock model and tokenizer.""" mock_model = MagicMock() mock_model.config = MagicMock() - mock_model.config.n_kv_heads = 8 - mock_model.config.n_heads = 8 - mock_model.config.dim = 128 - mock_model.config.n_layers = 2 - mock_model.config.max_len = 100 + mock_model.config.num_key_value_heads = 8 + mock_model.config.num_attention_heads = 8 + mock_model.config.hidden_size = 128 + mock_model.config.num_hidden_layers = 2 + mock_model.config.max_position_embeddings = 100 mock_model.parameters.return_value = iter( [MagicMock(dtype=torch.float32, device=torch.device("cpu"))] ) @@ -213,13 +213,13 @@ def _make_real_scheduler(device): cfg = AutoRegressiveLMConfig( vocab_size=200, - dim=16, - n_heads=2, - n_kv_heads=1, - dim_ffn=32, - max_len=64, - n_layers=2, - norm_eps=1e-5, + hidden_size=16, + num_attention_heads=2, + num_key_value_heads=1, + intermediate_size=32, + max_position_embeddings=64, + num_hidden_layers=2, + rms_norm_eps=1e-5, ) model = AutoRegressiveLM(cfg).to(device=device).eval() tokenizer = _Tok() diff --git a/tests/module/test_encoder.py b/tests/module/test_encoder.py index b66a595..1af7b85 100644 --- a/tests/module/test_encoder.py +++ b/tests/module/test_encoder.py @@ -12,13 +12,13 @@ from astrai.model.encoder import EmbeddingEncoder TINY_CONFIG = dict( vocab_size=128, - dim=8, - n_heads=2, - n_kv_heads=1, - dim_ffn=16, - max_len=64, - n_layers=2, - norm_eps=1e-5, + hidden_size=8, + num_attention_heads=2, + num_key_value_heads=1, + intermediate_size=16, + max_position_embeddings=64, + num_hidden_layers=2, + rms_norm_eps=1e-5, ) _device = "cuda" if torch.cuda.is_available() else "cpu" @@ -42,7 +42,7 @@ def test_encoder_forward_pooling(pooling_type): with torch.no_grad(): output = model(input_ids) - assert output.shape == (batch_size, TINY_CONFIG["dim"]) + assert output.shape == (batch_size, TINY_CONFIG["hidden_size"]) assert not torch.isnan(output).any() @@ -60,7 +60,7 @@ def test_encoder_forward_with_padding(): with torch.no_grad(): output = model(input_ids, input_mask=input_mask) - assert output.shape == (batch_size, TINY_CONFIG["dim"]) + assert output.shape == (batch_size, TINY_CONFIG["hidden_size"]) assert not torch.isnan(output).any() @@ -90,7 +90,7 @@ def test_encoder_from_transformer_checkpoint(): model = _make_model() state_dict = model.state_dict() state_dict["lm_head.weight"] = torch.randn( - TINY_CONFIG["vocab_size"], TINY_CONFIG["dim"], device=_device + TINY_CONFIG["vocab_size"], TINY_CONFIG["hidden_size"], device=_device ) new_model = _make_model() diff --git a/tests/module/test_forward_configs.py b/tests/module/test_forward_configs.py index 1662506..da00c32 100644 --- a/tests/module/test_forward_configs.py +++ b/tests/module/test_forward_configs.py @@ -6,13 +6,13 @@ from astrai.model.transformer import AutoRegressiveLM TINY_CONFIG = dict( vocab_size=128, - dim=8, - n_heads=2, - n_kv_heads=1, - dim_ffn=16, - max_len=64, - n_layers=2, - norm_eps=1e-5, + hidden_size=8, + num_attention_heads=2, + num_key_value_heads=1, + intermediate_size=16, + max_position_embeddings=64, + num_hidden_layers=2, + rms_norm_eps=1e-5, ) @@ -58,8 +58,13 @@ CONFIGS = [ id="gqa_qk_norm", ), pytest.param( - {**TINY_CONFIG, "attn_type": "gqa", "ffn_type": "mlp", "tie_weight": True}, - id="gqa_tie_weight", + { + **TINY_CONFIG, + "attn_type": "gqa", + "ffn_type": "mlp", + "tie_word_embeddings": True, + }, + id="gqa_tie_word_embeddings", ), ] @@ -82,7 +87,11 @@ def test_model_forward(config_kwargs): assert "logits" in output assert "hidden_states" in output assert output["logits"].shape == (batch_size, seq_len, config.vocab_size) - assert output["hidden_states"].shape == (batch_size, seq_len, config.dim) + assert output["hidden_states"].shape == ( + batch_size, + seq_len, + config.hidden_size, + ) assert not torch.isnan(output["logits"]).any() assert not torch.isnan(output["hidden_states"]).any() diff --git a/tests/module/test_lora.py b/tests/module/test_lora.py index d5b32b5..0a89a41 100644 --- a/tests/module/test_lora.py +++ b/tests/module/test_lora.py @@ -19,13 +19,13 @@ from astrai.model.components.lora import ( MODEL_KWARGS = dict( vocab_size=1000, - dim=64, - n_heads=4, - n_kv_heads=2, - dim_ffn=128, - n_layers=2, - max_len=32, - norm_eps=1e-5, + hidden_size=64, + num_attention_heads=4, + num_key_value_heads=2, + intermediate_size=128, + num_hidden_layers=2, + max_position_embeddings=32, + rms_norm_eps=1e-5, ) @@ -192,7 +192,7 @@ def test_inject_lora_on_moe_model(): n_routed_experts=4, n_shared_experts=1, n_activated_experts=2, - dim_ffn=32, + intermediate_size=32, ) inject_lora(model, r=4, alpha=8, target_modules={"up", "gate", "down"}) assert _get_lora_count(model) > 0 diff --git a/tests/module/test_tie_weight.py b/tests/module/test_tie_weight.py index f091abb..d60ed7b 100644 --- a/tests/module/test_tie_weight.py +++ b/tests/module/test_tie_weight.py @@ -17,13 +17,13 @@ def transformer_test_env(): config = { "vocab_size": 1000, - "dim": 8, - "n_heads": 2, - "n_kv_heads": 1, - "dim_ffn": 16, - "max_len": 64, - "n_layers": 2, - "norm_eps": 1e-5, + "hidden_size": 8, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "intermediate_size": 16, + "max_position_embeddings": 64, + "num_hidden_layers": 2, + "rms_norm_eps": 1e-5, } with open(config_path, "w") as f: @@ -45,7 +45,7 @@ def test_tie_weight_init(transformer_test_env): config_data = transformer_test_env["config"].copy() # case 1: tie weight - config_data["tie_weight"] = True + config_data["tie_word_embeddings"] = True with open(config_path, "w") as f: json.dump(config_data, f) @@ -63,7 +63,7 @@ def test_tie_weight_init(transformer_test_env): assert not torch.equal(model.lm_head.weight, original_weight) # case 2: not tie weight - config_data["tie_weight"] = False + config_data["tie_word_embeddings"] = False with open(config_path, "w") as f: json.dump(config_data, f) @@ -88,7 +88,7 @@ def test_model_save_load_with_tie_weight(transformer_test_env): config_data = transformer_test_env["config"].copy() # case 1: tie weight - config_data["tie_weight"] = True + config_data["tie_word_embeddings"] = True config_path = os.path.join(test_dir, "config.json") with open(config_path, "w") as f: @@ -108,7 +108,7 @@ def test_model_save_load_with_tie_weight(transformer_test_env): assert "lm_head.weight" not in model.state_dict() # case 2: not tie weight (form tie-weight state dict load) - config_data["tie_weight"] = False + config_data["tie_word_embeddings"] = False with open(config_path, "w") as f: json.dump(config_data, f) diff --git a/tests/trainer/test_grpo_strategy.py b/tests/trainer/test_grpo_strategy.py index 99e683a..856d832 100644 --- a/tests/trainer/test_grpo_strategy.py +++ b/tests/trainer/test_grpo_strategy.py @@ -13,16 +13,16 @@ class _FakeExecutor: return model.state_dict() -def _make_config(vocab_size=200, max_len=64): +def _make_config(vocab_size=200, max_position_embeddings=64): return AutoRegressiveLMConfig( vocab_size=vocab_size, - dim=16, - n_heads=2, - n_kv_heads=1, - dim_ffn=32, - max_len=max_len, - n_layers=2, - norm_eps=1e-5, + hidden_size=16, + num_attention_heads=2, + num_key_value_heads=1, + intermediate_size=32, + max_position_embeddings=max_position_embeddings, + num_hidden_layers=2, + rms_norm_eps=1e-5, ) diff --git a/tests/trainer/test_online_strategy.py b/tests/trainer/test_online_strategy.py index c8f78c7..f299a56 100644 --- a/tests/trainer/test_online_strategy.py +++ b/tests/trainer/test_online_strategy.py @@ -33,16 +33,16 @@ class _FakeExecutor: return model.state_dict() -def _make_config(vocab_size=200, max_len=64): +def _make_config(vocab_size=200, max_position_embeddings=64): return AutoRegressiveLMConfig( vocab_size=vocab_size, - dim=16, - n_heads=2, - n_kv_heads=1, - dim_ffn=32, - max_len=max_len, - n_layers=2, - norm_eps=1e-5, + hidden_size=16, + num_attention_heads=2, + num_key_value_heads=1, + intermediate_size=32, + max_position_embeddings=max_position_embeddings, + num_hidden_layers=2, + rms_norm_eps=1e-5, ) diff --git a/tests/trainer/test_rollout.py b/tests/trainer/test_rollout.py index 05677f0..3c73df4 100644 --- a/tests/trainer/test_rollout.py +++ b/tests/trainer/test_rollout.py @@ -71,16 +71,16 @@ class ConstantRewardModel(BaseRewardModel): return torch.full((B, G), float(self.value)) -def _make_config(vocab_size=200, max_len=128): +def _make_config(vocab_size=200, max_position_embeddings=128): return AutoRegressiveLMConfig( vocab_size=vocab_size, - dim=16, - n_heads=2, - n_kv_heads=1, - dim_ffn=32, - max_len=max_len, - n_layers=2, - norm_eps=1e-5, + hidden_size=16, + num_attention_heads=2, + num_key_value_heads=1, + intermediate_size=32, + max_position_embeddings=max_position_embeddings, + num_hidden_layers=2, + rms_norm_eps=1e-5, ) @@ -158,7 +158,7 @@ def _make_generator(device, **kw): model, tokenizer, max_batch_size=kw.get("max_batch_size", 8), - max_len=kw.get("max_len", 128), + max_len=kw.get("max_position_embeddings", 128), ) generator = RolloutGenerator( scheduler=scheduler, @@ -254,7 +254,7 @@ def _make_runner(device, **kw): group_size=kw.get("group_size", 2), max_tokens=kw.get("max_tokens", 8), max_batch_size=kw.get("max_batch_size", 8), - max_len=kw.get("max_len", 128), + max_len=kw.get("max_position_embeddings", 128), ) rm = ConstantRewardModel(1.0) return (