ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor : align config field names with Hugging Face

Browse Source 
- dim -> hidden_size, n_layers -> num_hidden_layers
- dim_ffn -> intermediate_size, n_heads -> num_attention_heads
- n_kv_heads -> num_key_value_heads, max_len -> max_position_embeddings
- norm_eps -> rms_norm_eps, tie_weight -> tie_word_embeddings
- update model, inference, training, scripts, tests, docs
This commit is contained in:
ViperEkura
2026-07-20 22:05:31 +08:00
parent d7ac66fb73
commit 0c86c89af4
23 changed files with 202 additions and 166 deletions
Download Patch File Download Diff File Expand all files Collapse all files
assets/docs/architecture.md
+15 -15
View File
@@ -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
assets/docs/dataflow.md
+1 -1
View File
@@ -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)
assets/docs/inference.md
+2 -2
View File
@@ -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
```
assets/docs/params.md
+2 -2
View File
@@ -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
astrai/config/model_config.py
+15 -15
View File
@@ -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
astrai/inference/core/scheduler.py
+6 -6
View File
@@ -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,
astrai/model/components/decoder_block.py
+11 -3
View File
@@ -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(
astrai/model/encoder.py
+13 -5
View File
@@ -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
astrai/model/transformer.py
+17 -9
View File
@@ -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]
astrai/trainer/train_context.py
+1 -1
View File
@@ -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,
scripts/eval/evaluate_mmlu.py
+1 -1
View File
@@ -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:]
scripts/tools/benchmark.py
+12 -12
View File
@@ -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(
scripts/tools/generate.py
+5 -2
View File
@@ -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",
scripts/tools/train.py
+1 -1
View File
@@ -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"),
tests/conftest.py
+14 -14
View File
@@ -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,
)
tests/inference/test_scheduler.py
+12 -12
View File
@@ -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()
tests/module/test_encoder.py
+10 -10
View File
@@ -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()
tests/module/test_forward_configs.py
+19 -10
View File
@@ -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()
tests/module/test_lora.py
+8 -8
View File
@@ -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
tests/module/test_tie_weight.py
+11 -11
View File
@@ -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)
tests/trainer/test_grpo_strategy.py
+8 -8
View File
@@ -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,
)
tests/trainer/test_online_strategy.py
+8 -8
View File
@@ -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,
)
tests/trainer/test_rollout.py
+10 -10
View File
@@ -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 (