amachine.am_transformers.am_models

  1try:
  2    from tokenizers import Tokenizer
  3    from tokenizers.models import BPE
  4    from tokenizers.pre_tokenizers import Split
  5    from transformers import (
  6        AutoTokenizer,
  7        AutoConfig,
  8        AutoModelForCausalLM,
  9        PreTrainedTokenizerFast, 
 10        GenerationConfig,
 11        GraniteMoeHybridConfig, 
 12        GraniteMoeHybridForCausalLM
 13    )
 14    _HAVE_TRANSFORMERS = True
 15except ImportError:
 16    _HAVE_TRANSFORMERS = False
 17
 18import warnings
 19
 20try:
 21    from .am_control_model_exp import *
 22except Exception:
 23    import traceback
 24    traceback.print_exc()
 25    warnings.warn("Failed to import control model")
 26
 27from typing import Literal
 28import copy
 29import json
 30import math
 31import os
 32from pathlib import Path
 33
 34# ---------------------------------------------------------------------------
 35# Validation
 36# ---------------------------------------------------------------------------
 37
 38def validate_model_behaviors(model_dir: str, verbose: bool = True) -> dict:
 39    """
 40    Sanity-check the decorator-driven behaviors that are easy to silently break
 41    when subclassing/overriding forward():
 42      1. AutoConfig/AutoModelForCausalLM registration resolves correctly
 43      2. use_cache defaults from config when not passed explicitly (@merge_with_config_defaults)
 44      3. return_dict=False returns a tuple, not a dataclass (@capture_outputs / @can_return_tuple)
 45      4. output_hidden_states / output_attentions actually populate (submodule hook matching)
 46      5. router_logits / aux_loss populate when the model has experts
 47
 48    Returns a dict of {check_name: True/False/None}. None means "not applicable"
 49    (e.g. router_logits check on a non-MoE model), not a failure.
 50
 51    Does not raise on failure by default — prints a report so you can eyeball
 52    which checks are genuine bugs vs. known/acceptable limitations (e.g. hidden_states
 53    capturing not firing because a custom layer class replaced the recorded one).
 54    """
 55    if not _HAVE_TRANSFORMERS:
 56        raise ImportError("validate_model_behaviors requires 'torch' and 'transformers'.")
 57
 58    import torch
 59    from transformers import AutoConfig, AutoModelForCausalLM
 60
 61    results: dict = {}
 62
 63    def report(name: str, ok, detail: str = ""):
 64        results[name] = ok
 65        if not verbose:
 66            return
 67        status = "PASS" if ok is True else "WARN/FAIL" if ok is False else "SKIP"
 68        print(f"  [{status}] {name}" + (f" — {detail}" if detail else ""))
 69
 70    print(f"\nValidating {model_dir}")
 71    print("-" * 60)
 72
 73    # --- 1. Registration / round trip ---------------------------------
 74    try:
 75        config = AutoConfig.from_pretrained(model_dir)
 76        model = AutoModelForCausalLM.from_pretrained(model_dir)
 77        model.eval()
 78        report(
 79            "registration_resolves",
 80            True,
 81            f"model_type={config.model_type!r} -> {type(model).__name__}",
 82        )
 83    except Exception as e:
 84        report("registration_resolves", False, f"{type(e).__name__}: {e}")
 85        print("  Cannot continue — registration must succeed before other checks.")
 86        return results
 87
 88    vocab_size = getattr(config, "vocab_size", 233)
 89    seq_len = 8
 90    input_ids = torch.randint(0, vocab_size, (1, seq_len))
 91
 92    # --- 2. use_cache defaults from config ------------------------------
 93    try:
 94        with torch.no_grad():
 95            out = model(input_ids)  # use_cache not passed explicitly
 96        expected = bool(getattr(config, "use_cache", True))
 97        got_cache = out.past_key_values is not None
 98        ok = got_cache == expected
 99        report(
100            "use_cache_defaults_from_config",
101            ok,
102            f"config.use_cache={expected}, past_key_values is not None: {got_cache}",
103        )
104    except Exception as e:
105        report("use_cache_defaults_from_config", False, f"{type(e).__name__}: {e}")
106
107    # --- 3. return_dict=False -> tuple ----------------------------------
108    try:
109        with torch.no_grad():
110            out = model(input_ids, return_dict=False)
111        ok = isinstance(out, tuple)
112        report("return_dict_false_gives_tuple", ok, f"got {type(out).__name__}")
113    except Exception as e:
114        report("return_dict_false_gives_tuple", False, f"{type(e).__name__}: {e}")
115
116    # --- 4. output_hidden_states / output_attentions populate -----------
117    try:
118        with torch.no_grad():
119            out = model(input_ids, output_hidden_states=True, output_attentions=True)
120        hs_ok = getattr(out, "hidden_states", None) is not None
121        attn_ok = getattr(out, "attentions", None) is not None
122        report(
123            "output_hidden_states_populates",
124            hs_ok,
125            "None — likely _can_record_outputs still points at the original layer "
126            "class, not your custom one" if not hs_ok else "",
127        )
128        report(
129            "output_attentions_populates",
130            attn_ok,
131            "None — same cause; attention-only layers (mamba layers have none) "
132            "may legitimately be absent" if not attn_ok else "",
133        )
134    except Exception as e:
135        report("output_hidden_states_populates", False, f"{type(e).__name__}: {e}")
136        report("output_attentions_populates", False, f"{type(e).__name__}: {e}")
137
138    # --- 5. router_logits / aux_loss for MoE models ----------------------
139    num_local_experts = getattr(config, "num_local_experts", 0)
140    if num_local_experts and num_local_experts > 0:
141        try:
142            labels = input_ids.clone()
143            with torch.no_grad():
144                out = model(
145                    input_ids,
146                    labels=labels,
147                    output_router_logits=True,
148                )
149            router_ok = getattr(out, "router_logits", None) is not None
150            aux_ok = getattr(out, "aux_loss", None) is not None
151            report("router_logits_populates", router_ok)
152            report("aux_loss_computed", aux_ok)
153        except Exception as e:
154            report("router_logits_populates", False, f"{type(e).__name__}: {e}")
155            report("aux_loss_computed", False, f"{type(e).__name__}: {e}")
156    else:
157        report("router_logits_populates", None, "skipped — num_local_experts=0")
158        report("aux_loss_computed", None, "skipped — num_local_experts=0")
159
160    print("-" * 60)
161    n_fail = sum(v is False for v in results.values())
162    print(f"{n_fail} check(s) need attention.\n" if n_fail else "All checks passed.\n")
163
164    return results
165
166
167BASE_CONFIG = {
168    "architectures": ["GraniteMoeHybridForCausalLM"],
169    "attention_bias": False,
170    "attention_dropout": 0.0,
171    "attention_multiplier": 0.015625,
172    "bos_token_id": 230,
173    "dtype": "float32",
174    "embedding_multiplier": 12,
175    "eos_token_id": 230,
176    "hidden_act": "silu",
177    "hidden_size": 128,
178    "init_method": "mup",
179    "initializer_range": 0.1,
180    "intermediate_size": 256,
181    "layer_types": ["attention", "attention", "attention"],
182    "logits_scaling": 1,
183    "mamba_chunk_size": 256,
184    "mamba_conv_bias": True,
185    "mamba_d_conv": 4,
186    "mamba_d_head": 16,
187    "mamba_d_state": 256,
188    "mamba_expand": 2,
189    "mamba_n_groups": 1,
190    "mamba_n_heads": 16,
191    "mamba_proj_bias": False,
192    "max_position_embeddings": 16384,
193    "model_type": "granitemoehybrid",
194    "normalization_function": "rmsnorm",
195    "num_attention_heads": 4,
196    "num_experts_per_tok": 0,
197    "num_hidden_layers": 3,
198    "num_key_value_heads": 1,
199    "num_local_experts": 0,
200    "output_router_logits": False,
201    "pad_token_id": 229, # Updated to match your dynamic tokenizer
202    "position_embedding_type": "rope",
203    "residual_multiplier": 0.22,
204    "rms_norm_eps": 1e-05,
205    "rope_scaling": None,
206    "rope_theta": 10000000,
207    "router_aux_loss_coef": 0.01,
208    "shared_intermediate_size": 256,
209    "tie_word_embeddings": True,
210    "use_cache": True,
211    "vocab_size": 233,
212}
213
214# ---------------------------------------------------------------------------
215# Helpers
216# ---------------------------------------------------------------------------
217
218def _derive_mamba_n_heads(width: int, mamba_expand: int, start: int = 16) -> int:
219    """
220    Find the largest power-of-2 <= start such that:
221      - (width * mamba_expand) is divisible by mamba_n_heads
222      - mamba_d_head = (width * mamba_expand) / mamba_n_heads >= 8
223    """
224    mamba_inner = width * mamba_expand
225    n = start
226    while n > 1 and (mamba_inner % n != 0 or mamba_inner // n < 8):
227        n //= 2
228    return n
229
230def _rope_theta_for_seq_len(seq_len: int, head_dim: int) -> int:
231    """Derive rope_theta from target sequence length and head dimension."""
232    return int(math.ceil(seq_len ** (head_dim / (head_dim - 2))))
233
234def _make_layer_types(num_layers: int, mamba_ratio: float) -> list[str]:
235    """
236    Build a layer_types list of length num_layers with mamba layers distributed
237    as evenly as possible among attention layers.
238    """
239    if not 0.0 <= mamba_ratio <= 1.0:
240        raise ValueError(f"mamba_ratio must be in [0, 1], got {mamba_ratio}")
241
242    num_mamba = round(num_layers * mamba_ratio)
243    types = ["attention"] * num_layers
244
245    for i in range(num_mamba):
246        pos = int((i + 0.5) * num_layers / num_mamba)
247        types[pos] = "mamba"
248
249    return types
250
251# ---------------------------------------------------------------------------
252# Config building
253# ---------------------------------------------------------------------------
254
255def build_config(
256    num_layers: int,
257    hidden_size: int,
258    ffn_multiplier: float = 2.0,
259    head_dim: int = 64,
260    num_kv_heads: int = 1,
261    logits_scaling: float | None = None,
262    mamba_n_heads: int = 16,
263    mamba_ratio: float = 0.0,
264    num_local_experts: int = 0,
265    num_experts_per_tok: int = 0,
266    position_embedding_type : str ="rope",
267    rope_theta: int = 1000000,
268    tie_embeddings: bool = False,
269    vocab_size: int = 233,
270    pad_token_id: int = 229,
271    bos_token_id: int = 230,
272    eos_token_id: int = 230,
273    model_type: str | None = None,
274    model_config_kwargs : dict[ str, str | bool | int | float ] | None = None,
275    dtype : Literal[ "float32", "bfloat16" ] = "float32"
276) -> dict:
277    """Build a config dict for a GraniteMoeHybrid model."""
278    if hidden_size % head_dim != 0:
279        raise ValueError(f"hidden_size ({hidden_size}) must be divisible by head_dim ({head_dim})")
280
281    num_attention_heads = hidden_size // head_dim
282
283    if num_attention_heads % num_kv_heads != 0:
284        raise ValueError(f"num_attention_heads ({num_attention_heads}) must be divisible by num_kv_heads ({num_kv_heads})")
285
286    if num_local_experts > 0 and num_experts_per_tok <= 0:
287        raise ValueError(f"num_experts_per_tok must be > 0 when num_local_experts ({num_local_experts}) > 0")
288
289    if num_experts_per_tok > num_local_experts:
290        raise ValueError(f"num_experts_per_tok ({num_experts_per_tok}) cannot exceed num_local_experts ({num_local_experts})")
291
292    attention_multiplier = 1.0 / head_dim
293    intermediate_size = int(hidden_size * ffn_multiplier)
294
295    mamba_expand = BASE_CONFIG["mamba_expand"]
296    mamba_inner_dim = hidden_size * mamba_expand
297    if mamba_inner_dim % mamba_n_heads != 0:
298        raise ValueError(f"hidden_size * mamba_expand ({mamba_inner_dim}) must be divisible by mamba_n_heads ({mamba_n_heads}).")
299    mamba_d_head = mamba_inner_dim // mamba_n_heads
300
301    if logits_scaling is None:
302        logits_scaling = hidden_size / 256
303
304    layer_types = _make_layer_types(num_layers, mamba_ratio)
305
306    cfg = copy.deepcopy(BASE_CONFIG)
307    cfg.update({
308        "hidden_size":          hidden_size,
309        "num_hidden_layers":    num_layers,
310        "layer_types":          layer_types,
311        "intermediate_size":    intermediate_size,
312        "shared_intermediate_size": intermediate_size,
313        "num_attention_heads":  num_attention_heads,
314        "num_key_value_heads":  num_kv_heads,
315        "mamba_n_heads":        mamba_n_heads,
316        "mamba_d_head":         mamba_d_head,
317        "attention_multiplier": attention_multiplier,
318        "logits_scaling":       logits_scaling,
319        "rope_theta":           rope_theta,
320        "position_embedding_type" : position_embedding_type,
321        "num_local_experts":    num_local_experts,
322        "num_experts_per_tok":  num_experts_per_tok,
323        "output_router_logits": num_local_experts > 0,
324        "tie_word_embeddings" : tie_embeddings,
325        "pad_token_id":         pad_token_id,
326        "bos_token_id":         bos_token_id,
327        "eos_token_id":         eos_token_id,
328        "vocab_size": vocab_size,
329        "dtype" :dtype
330    })
331
332    if model_type is not None:
333        cfg["model_type"] = model_type
334
335    if model_config_kwargs :
336        cfg.update( model_config_kwargs )
337
338    return cfg
339
340# ---------------------------------------------------------------------------
341# Parameter estimation
342# ---------------------------------------------------------------------------
343
344def estimate_params(cfg: dict) -> int:
345    """Rough non-embedding parameter count for quick sanity checks."""
346    
347    H = cfg["hidden_size"]
348    I = cfg["intermediate_size"]
349    
350    layer_types       = cfg["layer_types"]
351    num_local_experts = cfg.get("num_local_experts", 0)
352    mamba_expand      = cfg.get("mamba_expand", 2)
353
354    attn_params = 4 * H * H
355    expert_count = max(1, num_local_experts)
356    ffn_params   = 3 * H * I * expert_count
357    mamba_params = 3 * H * (H * mamba_expand)
358    ln_params = 2 * H
359
360    total = 0
361    for lt in layer_types:
362        if lt == "mamba":
363            total += mamba_params + ln_params
364        else:
365            total += attn_params + ffn_params + ln_params
366
367    return total
368
369# ---------------------------------------------------------------------------
370# Saving
371# ---------------------------------------------------------------------------
372
373def save_config(cfg: dict, output_dir: str) -> str:
374    """Write config.json to output_dir and return the full path."""
375    os.makedirs(output_dir, exist_ok=True)
376    config_path = os.path.join(output_dir, "config.json")
377    with open(config_path, "w") as f:
378        json.dump(cfg, f, indent=2)
379    return config_path
380
381def save_model(cfg: dict, output_dir: str) -> None:
382    """
383    Instantiate a GraniteMoeHybridForCausalLM from cfg, save weights and config
384    to output_dir.
385    """
386    if not _HAVE_TRANSFORMERS:
387        raise ImportError(
388            "The 'save_model' function requires 'torch' and 'transformers'."
389        )
390
391    config_path = save_config(cfg, output_dir)
392
393    from pprint import pprint as pprint
394    try:
395        hf_config = AutoConfig.from_pretrained(output_dir)
396    except ValueError as e:
397        raise ValueError( f"Could not resolve model_type={cfg.get('model_type')!r}." ) from e
398
399    model = AutoModelForCausalLM.from_config(hf_config)
400
401    total_params = sum(p.numel() for p in model.parameters())
402    print(f"Actual total parameters: {total_params / 1e6:.2f}M")
403
404    model.save_pretrained(output_dir, safe_serialization=True)
405    print(f"Model saved to {output_dir}/")
406
407# ---------------------------------------------------------------------------
408# Ensemble generation
409# ---------------------------------------------------------------------------
410
411def _model_name( 
412    width: int, 
413    depth: int, 
414    head_dim : int,
415    mamba_ratio: float, 
416    num_local_experts: int,
417    position_embedding_type : str,
418    opt : str | None = None,
419    id_str : str | None = None,
420    model_type : str | None = None,
421    model_config_kwargs : dict[ str, str | bool | int | float ] | None = None ) -> str :
422    
423    mt = "g4-zepto" if model_type is None else model_type
424
425    parts = [f"{mt}-w{width}-d{depth}-h{head_dim}-p{position_embedding_type}"]
426    
427    if mamba_ratio > 0.0:
428        parts.append(f"mamba{mamba_ratio:.2f}".rstrip("0").rstrip("."))
429    if num_local_experts > 0:
430        parts.append(f"moe{num_local_experts}")
431    if opt is not None :
432        parts.append(f"o{opt}")
433    if id_str is not None :
434        parts.append(f"id{id_str}")
435    
436    return "-".join(parts)
437
438def get_text_tokenizer(model_name_or_path: str = "ibm-granite/granite-4.0-350m"):
439    
440    """Loads a pre-trained tokenizer for standard text modeling or a local tokenizer.json."""
441
442    if not _HAVE_TRANSFORMERS:
443        raise ImportError("Requires 'transformers'.")
444
445    is_local_json = os.path.isfile(model_name_or_path) and model_name_or_path.endswith(".json")
446
447    if is_local_json:
448        
449        raw_tokenizer = Tokenizer.from_file(model_name_or_path)
450        
451        def extract_token_str(token_val):
452            """Helper to handle Hugging Face tokenizer config dicts."""
453            if isinstance(token_val, dict):
454                return token_val.get("content", None)
455            return token_val
456
457        unk_token = None
458        pad_token = None
459        eos_token = None
460
461        hf_tk_path = Path( model_name_or_path ).parent / "tokenizer_config.json"
462
463        if hf_tk_path.exists():
464            with open( hf_tk_path, "r") as f:
465                tokenizer_config = json.load(f)
466                unk_token = extract_token_str(tokenizer_config.get('unk_token', None))
467                pad_token = extract_token_str(tokenizer_config.get('pad_token', None))
468                eos_token = extract_token_str(tokenizer_config.get('eos_token', None))
469
470        hf_tokenizer = PreTrainedTokenizerFast(
471            tokenizer_object=raw_tokenizer,
472            unk_token=unk_token,
473            eos_token=eos_token,
474            pad_token=pad_token
475        )
476    else:
477
478        hf_tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
479        
480        # Pre-trained models sometimes lack a defined pad token
481        assert hf_tokenizer.pad_token is not None
482            
483        raw_tokenizer = hf_tokenizer.backend_tokenizer
484        
485    vocabulary = hf_tokenizer.get_vocab()
486    
487    # 3. Safely build GenerationConfig
488    # Local tokenizers won't have a full model config, so we conditionally build kwargs
489    gen_kwargs = {
490        "eos_token_id": hf_tokenizer.eos_token_id,
491        "pad_token_id": hf_tokenizer.pad_token_id,
492    }
493    
494    if hf_tokenizer.bos_token_id is not None:
495        gen_kwargs["bos_token_id"] = hf_tokenizer.bos_token_id
496        
497    if not is_local_json:
498        gen_kwargs["_from_model_config"] = True
499        
500    gen_config = GenerationConfig(**gen_kwargs)
501
502    return raw_tokenizer, vocabulary, hf_tokenizer, gen_config
503
504def get_tokenizer():
505    
506    """Generates and saves a character-level BPE tokenizer (no merges)."""
507
508    if not _HAVE_TRANSFORMERS:
509        raise ImportError(
510            "The 'save_model' function requires 'torch' and 'transformers'. "
511        )
512
513    ascii_symbols = [chr(i) for i in range(32, 127)]
514    greek_upper = [chr(i) for i in range(0x0391, 0x03A5) if i != 0x03A2][:19]
515    greek_lower = [chr(i) for i in range(0x03B1, 0x03C5) if i != 0x03C2][:19]
516    geometric = [chr(i) for i in range(0x25A0, 0x2600)]
517
518    symbols = ascii_symbols + greek_upper + greek_lower + geometric
519    
520    assert len(symbols) == 229, f"Expected 229 content symbols, got {len(symbols)}"
521    assert len(set(symbols)) == len(symbols), "Duplicate symbol detected!"
522
523    special_tokens = ["<|pad|>", "<|end_of_text|>", "<|unk|>", "<|mask|>"]
524
525    PAD_ID          = 229
526    END_OF_TEXT_ID  = 230
527    UNK_ID          = 231
528    MASK_ID         = 232
529
530    vocab = {s: i for i, s in enumerate(symbols)}
531    vocab["<|pad|>"]          = PAD_ID
532    vocab["<|end_of_text|>"]  = END_OF_TEXT_ID
533    vocab["<|unk|>"]          = UNK_ID
534    vocab["<|mask|>"]         = MASK_ID
535
536    raw_tokenizer = Tokenizer(BPE(
537        unk_token="<|unk|>",
538        end_of_word_suffix="",
539        continuing_subword_prefix="",
540        merges=[],
541        vocab=vocab,
542    ))
543
544    raw_tokenizer.add_special_tokens(special_tokens)
545    raw_tokenizer.pre_tokenizer = Split(pattern="", behavior="isolated")
546
547    hf_tokenizer = PreTrainedTokenizerFast(
548        tokenizer_object=raw_tokenizer,
549        bos_token="<|end_of_text|>",
550        eos_token="<|end_of_text|>",
551        pad_token="<|pad|>",
552        unk_token="<|unk|>",
553        mask_token="<|mask|>",
554        clean_up_tokenization_spaces=False,
555        padding_side="left",
556    )
557
558    gen_config = GenerationConfig(
559        bos_token_id=hf_tokenizer.bos_token_id,
560        eos_token_id=hf_tokenizer.eos_token_id,
561        pad_token_id=hf_tokenizer.pad_token_id,
562        _from_model_config=True,
563    )
564
565    return raw_tokenizer, vocab, hf_tokenizer, gen_config
566
567def generate_ensemble(
568    output_dir: str,
569    core_params: list[dict],
570    dtype : Literal[ "float32", "bfloat16" ] = "float32",
571    seq_len: int = 1024,
572    tie_embeddings : bool = False,
573    mode : Literal["create", "extend"] = "create",
574    tokenizer_type : Literal[ "symbolic", "text" ] = "symbolic",
575    tokenizer_id : str = "ibm-granite/granite-4.0-350m"
576) -> dict:
577    """Generate a sweep of toy models with varying widths, depths, and architectures."""
578
579    if not _HAVE_TRANSFORMERS:
580        raise ImportError(
581            "The 'save_model' function requires 'torch' and 'transformers'. "
582        )
583
584    output_path = Path(output_dir)
585    output_path.mkdir(parents=True, exist_ok=True)
586
587    manifest = {
588        "models": [],
589    }
590
591    # Generate the default simple character level tokenizer used for HMM data
592    if tokenizer_type == "symbolic" :
593        raw_tokenizer, vocabulary, hf_tokenizer, gen_config = get_tokenizer()
594    elif tokenizer_type == "text" :
595        raw_tokenizer, vocabulary, hf_tokenizer, gen_config = get_text_tokenizer( tokenizer_id )
596    else :
597        raise ValueError( "Unsupported tokenizer_type {tokenizer_type}. Options are symbolic or text" )
598
599    vocab_size = len(hf_tokenizer)
600
601    pad_token_id : int = hf_tokenizer.pad_token_id
602    bos_token_id : int = hf_tokenizer.bos_token_id
603    eos_token_id : int = hf_tokenizer.eos_token_id
604
605    for idx, p in enumerate( core_params ):
606
607        if not "width" in p or not "depth" in p or not "head_dim" in p :
608            raise ValueError( "core params must contain width, depth, and head_dim" )
609
610        width    = int(p["width"])
611        depth    = int(p["depth"])
612        head_dim = int(p["head_dim"])
613
614        use_rope = p.get( "use_rope", True )
615
616        opt    = p.get( "opt", None )
617        id_str = p.get( "id", str(idx) )
618
619        position_embedding_type = "rope" if use_rope else "nope"
620        mamba_ratio             = float(p.get("mamba_ratio", 0.0))
621        num_local_experts       = int(p.get("num_local_experts", 0))
622        num_experts_per_tok     = int(p.get("num_experts_per_tok", 0))
623        seq_len                 = int(p.get("seq_len", 1024))
624        ffn_multiplier          = float(p.get("ffn_multiplier", 2.0))
625
626        rope_theta          = _rope_theta_for_seq_len(seq_len, head_dim)
627        num_attention_heads = width // head_dim
628        kv_heads            = max(1, num_attention_heads)
629        mamba_n_heads       = _derive_mamba_n_heads(width, BASE_CONFIG["mamba_expand"])
630
631        model_type   = p.get("model_type", None)
632        model_config_kwargs = p.get("model_config_kwargs", None)
633
634        name = _model_name(
635            width=width, 
636            depth=depth, 
637            head_dim=head_dim, 
638            mamba_ratio=mamba_ratio, 
639            num_local_experts=num_local_experts,
640            position_embedding_type=position_embedding_type,
641            opt=opt,
642            id_str=id_str,
643            model_type=model_type,
644            model_config_kwargs=model_config_kwargs
645        )
646
647        model_dir = str(output_path / name)
648        model_path = Path(model_dir)
649
650        layer_types = _make_layer_types(depth, mamba_ratio)
651        num_mamba_layers = layer_types.count("mamba")
652        num_attn_layers  = layer_types.count("attention")
653
654        print(f"\n{'=' * 60}")
655        print(f"Generating {name}")
656        print(f"  width={width}, depth={depth}, "
657              f"attn_layers={num_attn_layers}, mamba_layers={num_mamba_layers}")
658        
659        if num_local_experts:
660            print(f"  MoE: {num_local_experts} experts, top-{num_experts_per_tok}")
661
662        cfg = build_config(
663            num_layers=depth,
664            hidden_size=width,
665            head_dim=head_dim,
666            ffn_multiplier=ffn_multiplier,
667            num_kv_heads=kv_heads,
668            mamba_n_heads=mamba_n_heads,
669            mamba_ratio=mamba_ratio,
670            num_local_experts=num_local_experts,
671            num_experts_per_tok=num_experts_per_tok,
672            position_embedding_type=position_embedding_type,
673            rope_theta=rope_theta,
674            tie_embeddings=tie_embeddings,
675            vocab_size=vocab_size,
676            pad_token_id=pad_token_id,
677            bos_token_id=bos_token_id,
678            eos_token_id=eos_token_id,
679            model_type=model_type,
680            model_config_kwargs=model_config_kwargs,
681            dtype=dtype
682        )
683
684        approx = estimate_params(cfg)
685        print(f"  Estimated non-embedding params: ~{approx / 1e6:.2f}M")
686
687        save_model(cfg, model_dir)
688
689        raw_tokenizer.save(str(model_path / "tokenizer.json"))
690        
691        with open(model_path / "vocab.json", "w", encoding="utf-8") as f:
692            json.dump(vocabulary, f, indent=2, ensure_ascii=False)
693
694        hf_tokenizer.save_pretrained(model_dir)
695        gen_config.save_pretrained(model_dir)
696
697        manifest["models"].append({
698            "name":                     name,
699            "hidden_size":              width,
700            "num_hidden_layers":        depth,
701            "num_attention_layers":     num_attn_layers,
702            "head_dim":                 head_dim,
703            "rope_theta" :              rope_theta,
704            "num_mamba_layers":         num_mamba_layers,
705            "mamba_ratio":              mamba_ratio,
706            "num_attention_heads":      num_attention_heads,
707            "num_key_value_heads":      kv_heads,
708            "intermediate_size":        cfg["intermediate_size"],
709            "mamba_n_heads":            mamba_n_heads,
710            "mamba_d_head":             cfg["mamba_d_head"],
711            "num_local_experts":        num_local_experts,
712            "num_experts_per_tok":      num_experts_per_tok,
713            "estimated_params":         approx,
714            "position_embedding_type" : position_embedding_type,
715            "seq_len" :                 seq_len
716        } )
717
718        validate_model_behaviors( model_path, True )
719
720    manifest_path = output_path / "manifest.json"
721    
722    with open(manifest_path, "w") as f:
723        json.dump(manifest, f, indent=2)
724    
725    print(f"\nManifest written to {manifest_path}")
726
727    return manifest