| --- |
| library_name: vllm |
| pipeline_tag: text-generation |
| base_model: |
| - nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-BF16 |
| --- |
| |
| This tiny model is randomly initialized for debugging using [nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-BF16](https://huggingface.co/nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-BF16). |
|
|
| Dory preserves two recurrent passes, input reinjection, nGPT residuals, |
| SWA/full attention, and untied BF16 token matrices. The pinned official vLLM |
| implementation is included with tensor-parallel vector loading and distributed |
| nGPT normalization. Transformers has no Dory model implementation or source |
| auto_map; checkpoint generation follows the official source tensor schema. |
| |
| | Feature | Source | Tiny | |
| |---|---|---| |
| | Checkpoint tensor elements | 3,975,229,440 BF16 | 689,248 BF16 | |
| | Safetensors weight files | 7,950,515,472 bytes (model.safetensors) | 1,412,480 bytes (model.safetensors) | |
| | Input / recurrent / output sublayers | 16 / 40 / 16 | 16 / 16 / 16 | |
| | Recurrent passes | 2 | 2 | |
| | SWA/full cycle | 3 SWA+MLP pairs, then full+MLP | Same; two cycles per partition | |
| | Vocabulary / embedding tying | 131,072 / false | Same | |
| | Hidden width / MLP width | 2,560 / 7,168 | 2 / 32 | |
| | Q heads / KV heads / head dimension | 16 / 4 / 256 | 4 / 1 / 32 | |
| | Checkpoint precision | BF16 | Same | |
| |
| Install the included runtime once with `pip install --no-deps <snapshot directory>` |
| before running the offline example. It registers Dory in every vLLM worker. |
| The native vLLM sampler avoids FlashInfer sampling JIT; attention uses FlashAttention. |
| Width 1 was rejected because scalar nGPT normalization collapses the latent state: |
| one and two recurrent passes produced identical final states. Width 2 retains |
| continuous recurrent updates; both untied vocabulary matrices keep their source rows. |
| TP=1 and TP=4 were validated with English, Chinese and a 1,100-token prompt, |
| generating ten tokens each. Both per-loop and last-loop caches and CUDA Graph |
| were exercised; TP=1/TP=4 returned identical token IDs with at most 0.000644 |
| logprob difference. The independent creation and offline snippets passed outside |
| the checkout with repository imports blocked. |
| The current offline and HTTP serving examples were verified on four NVIDIA H20 |
| GPUs with TP=4, FlashAttention and CUDA Graph. Both chat requests (English and |
| Chinese) and a 1,100-token completion request generated ten tokens with finite |
| log probabilities. |
| |
| | File path | Size | |
| |------|------| |
| | model.safetensors | 1.4MB | |
| |
| ### Example usage: |
| |
| - vLLM offline (✅ verified) |
| |
| ```python |
| def run_vllm(model_id: str): |
| import os |
| from pathlib import Path |
| from huggingface_hub import snapshot_download |
| from transformers import AutoTokenizer |
| |
| # Use vLLM's native sampler; FlashAttention and CUDA Graph remain enabled. |
| os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0") |
| os.environ.setdefault("VLLM_ALLREDUCE_USE_FLASHINFER", "0") |
| os.environ.setdefault("VLLM_ALLREDUCE_USE_SYMM_MEM", "0") |
| from vllm import LLM, SamplingParams |
| |
| model_id = str(model_id) |
| folder = Path(model_id) if Path(model_id).is_dir() else Path(snapshot_download(model_id)) |
| tokenizer = AutoTokenizer.from_pretrained(folder, trust_remote_code=True) |
| prompt = tokenizer.apply_chat_template( |
| [{"role": "user", "content": "Hello"}], tokenize=False, add_generation_prompt=True |
| ) |
| llm = LLM( |
| model=str(folder), |
| trust_remote_code=True, |
| tensor_parallel_size=4, |
| max_model_len=2048, |
| max_num_batched_tokens=256, |
| max_num_seqs=4, |
| gpu_memory_utilization=0.2, |
| disable_custom_all_reduce=True, |
| compilation_config={"cudagraph_capture_sizes": [1, 2, 4, 8]}, |
| ) |
| try: |
| print(llm) |
| print(llm.generate([prompt], SamplingParams(temperature=0, max_tokens=10, ignore_eos=True))[0].outputs[0].text) |
| finally: |
| llm.llm_engine.engine_core.shutdown(timeout=30) |
| |
|
|
| if __name__ == "__main__": |
| run_vllm(model_id="tiny-random/dory-1.0") |
| ``` |
| |
| - vLLM serving (✅ verified) |
|
|
| Download `tiny-random/dory-1.0` and install its included runtime before starting the server. |
|
|
| ```bash |
| export VLLM_USE_FLASHINFER_SAMPLER=0 VLLM_ALLREDUCE_USE_FLASHINFER=0 VLLM_ALLREDUCE_USE_SYMM_MEM=0 |
| python - <<'PY' |
| from huggingface_hub import snapshot_download |
| snapshot_download("tiny-random/dory-1.0", local_dir="dory-1.0") |
| PY |
| pip install --no-deps ./dory-1.0 |
| |
| vllm serve ./dory-1.0 \ |
| --trust-remote-code \ |
| --tensor-parallel-size 4 \ |
| --hf-overrides '{"n_recurrent_loops": 2, "recurrent_kv_cache_mode": "per_loop"}' \ |
| --max-model-len 2048 \ |
| --max-num-batched-tokens 256 \ |
| --max-num-seqs 4 \ |
| --gpu-memory-utilization 0.2 \ |
| --disable-custom-all-reduce \ |
| --compilation-config '{"cudagraph_capture_sizes": [1, 2, 4, 8]}' |
| ``` |
|
|
| ### Codes to create this repo: |
|
|
| <details><summary>Click to expand</summary> |
|
|
| ```python |
| def create_model(source_model_id: str, source_revision: str, save_folder: str): |
| import json |
| import urllib.request |
| from pathlib import Path |
| |
| import torch |
| from huggingface_hub import hf_hub_download |
| from safetensors.torch import save_file |
| |
| runtime_revision = "b5007e42b09c3990685fe564a31ae6917394a80b" |
| folder = Path(save_folder) |
| folder.mkdir(parents=True, exist_ok=True) |
| config = json.loads(Path(hf_hub_download(source_model_id, "config.json", revision=source_revision)).read_text()) |
| # Two complete SWA/SWA/SWA/full cycles in every physical partition. |
| config.update( |
| hidden_size=2, |
| intermediate_size=32, |
| num_attention_heads=4, |
| num_key_value_heads=1, |
| head_dim=32, |
| num_hidden_layers=48, |
| n_input_layers=16, |
| n_recurrent_layers=16, |
| n_output_layers=16, |
| ) |
| config["layer_types"] = config["layer_types"][:48] |
| config["auto_map"] = {"AutoConfig": "configuration_dory.DoryConfig"} |
| (folder / "config.json").write_text(json.dumps(config, indent=2) + "\n") |
| for name in ( |
| "tokenizer.json", |
| "tokenizer_config.json", |
| "special_tokens_map.json", |
| "chat_template.jinja", |
| "generation_config.json", |
| ): |
| (folder / name).write_bytes(Path(hf_hub_download(source_model_id, name, revision=source_revision)).read_bytes()) |
| |
| # No Transformers model class or source auto_map exists. Preserve the pinned |
| # official checkpoint schema and ship its vLLM runtime with TP adaptations. |
| base = f"https://raw.githubusercontent.com/vllm-project/vllm/{runtime_revision}/" |
| for remote, local in ( |
| ("vllm/transformers_utils/configs/dory.py", "configuration_dory.py"), |
| ("vllm/model_executor/models/dory.py", "vllm_dory.py"), |
| ): |
| with urllib.request.urlopen(base + remote) as response: |
| code = response.read().decode() |
| if local == "vllm_dory.py": |
| code = code.replace( |
| "from .utils import maybe_prefix", "from vllm.model_executor.models.utils import maybe_prefix" |
| ) |
| code = code.replace( |
| "from vllm.distributed import get_tensor_model_parallel_world_size", |
| "from vllm.distributed import (get_tensor_model_parallel_world_size, " |
| "get_tensor_model_parallel_rank, tensor_model_parallel_all_reduce)", |
| ) |
| code = code.replace( |
| "EPS = 1e-6", |
| "EPS = 1e-6\n\n" |
| "def partition_norm(x):\n" |
| " total = tensor_model_parallel_all_reduce(x.float().square().sum(dim=-1, keepdim=True))\n" |
| " return (x.float() * torch.rsqrt(total + EPS)).to(x.dtype)\n\n" |
| "def load_vector(param, weight):\n" |
| " rank = get_tensor_model_parallel_rank()\n" |
| " param.data.copy_(weight.narrow(0, rank * param.numel(), param.numel()))\n", |
| ) |
| code = code.replace( |
| "self.num_heads = config.num_attention_heads", |
| "self.num_heads = config.num_attention_heads // get_tensor_model_parallel_world_size()", |
| ) |
| code = code.replace( |
| "self.num_kv_heads = config.num_key_value_heads", |
| "self.num_kv_heads = max(1, config.num_key_value_heads // get_tensor_model_parallel_world_size())", |
| ) |
| code = code.replace( |
| " self.num_heads,\n self.num_kv_heads,", |
| " config.num_attention_heads,\n config.num_key_value_heads,", |
| ) |
| code = code.replace( |
| " self.q_size,\n bias=False,", |
| " config.num_attention_heads * self.head_dim,\n bias=False,", |
| ) |
| code = code.replace( |
| " self.q_size,\n self.hidden_size,", |
| " config.num_attention_heads * self.head_dim,\n self.hidden_size,", |
| ) |
| code = code.replace( |
| "self.s_gate = nn.Parameter(torch.empty(self.q_size))", |
| "self.s_gate = nn.Parameter(torch.empty(self.q_size))\n" |
| " self.sqk.weight_loader = self.s_gate.weight_loader = load_vector", |
| ) |
| code = code.replace( |
| "self.ffn = config.intermediate_size", |
| "self.ffn = config.intermediate_size // get_tensor_model_parallel_world_size()", |
| ) |
| code = code.replace("[self.ffn] * 2", "[config.intermediate_size] * 2") |
| code = code.replace( |
| " self.ffn,\n config.hidden_size,", |
| " config.intermediate_size,\n config.hidden_size,", |
| ) |
| code = code.replace( |
| "self.suv_up = nn.Parameter(torch.empty(self.ffn))", |
| "self.suv_up = nn.Parameter(torch.empty(self.ffn))\n" |
| " self.suv_gate.weight_loader = self.suv_up.weight_loader = load_vector", |
| ) |
| code = code.replace("out = justnorm(out)", "out = partition_norm(out)") |
| code = code.replace("m = justnorm(m)", "m = partition_norm(m)") |
| code = code.replace( |
| " if get_tensor_model_parallel_world_size() != 1:\n" |
| ' raise ValueError("Dory currently supports tensor_parallel_size=1 only")\n', |
| "", |
| ) |
| code = code.replace( |
| " self._check_weight_shape(name, w, shape)\n" " getattr(param,", |
| ' if name.endswith((".sqk", ".s_gate")):\n' |
| " shape = (self.config.num_attention_heads * self.config.head_dim,)\n" |
| ' elif name.endswith((".suv_gate", ".suv_up")):\n' |
| " shape = (self.config.intermediate_size,)\n" |
| ' elif name.endswith(".g_proj.weight"):\n' |
| " shape = (self.config.num_attention_heads * self.config.head_dim, self.config.hidden_size)\n" |
| ' elif name.endswith(".o_proj.weight"):\n' |
| " shape = (self.config.hidden_size, self.config.num_attention_heads * self.config.head_dim)\n" |
| ' elif name.endswith(".down_proj.weight"):\n' |
| " shape = (self.config.hidden_size, self.config.intermediate_size)\n" |
| " self._check_weight_shape(name, w, shape)\n" |
| " getattr(param,", |
| ) |
| assert "tensor_parallel_size=1 only" not in code |
| assert code.count("partition_norm(") == 3 |
| (folder / local).write_text(code) |
| # vLLM spawns fresh interpreters. A general plugin registers the model in |
| # every process; parent-process ModelRegistry changes alone do not survive. |
| (folder / "dory_plugin.py").write_text( |
| "def register():\n" |
| " from vllm import ModelRegistry\n" |
| " ModelRegistry.register_model('DoryForCausalLM', 'vllm_dory:DoryForCausalLM')\n" |
| ) |
| (folder / "pyproject.toml").write_text( |
| '[build-system]\nrequires = ["setuptools>=61"]\nbuild-backend = "setuptools.build_meta"\n' |
| '[project]\nname = "dory-inference"\nversion = "0.1.0"\n' |
| '[project.entry-points."vllm.general_plugins"]\ndory = "dory_plugin:register"\n' |
| '[tool.setuptools]\npy-modules = ["dory_plugin", "vllm_dory", "configuration_dory"]\n' |
| ) |
| |
| torch.manual_seed(42) |
| h, f = config["hidden_size"], config["intermediate_size"] |
| q, kv = config["num_attention_heads"] * config["head_dim"], config["num_key_value_heads"] * config["head_dim"] |
| tensors = {} |
| |
| class CheckpointModule(torch.nn.Module): |
| def extra_repr(self): |
| return "\n".join( |
| f"{name}: shape={tuple(p.shape)}, dtype={p.dtype}" for name, p in self.named_parameters(recurse=False) |
| ) |
| |
| graph = CheckpointModule() |
| |
| def add(name, shape): |
| tensor = torch.randn(shape, dtype=torch.bfloat16) * config["initializer_range"] |
| tensors[name] = tensor |
| module = graph |
| parts = name.split(".") |
| for part in parts[:-1]: |
| if part not in module._modules: |
| module.add_module(part, CheckpointModule()) |
| module = module._modules[part] |
| module.register_parameter(parts[-1], torch.nn.Parameter(tensor)) |
| |
| add("backbone.embeddings.weight", (config["vocab_size"], h)) |
| add("lm_head.weight", (config["vocab_size"], h)) |
| add("lm_head.logit_scale", (config["vocab_size"],)) |
| for i, kind in enumerate(config["layer_types"]): |
| prefix = f"backbone.layers.{i}" |
| if kind == "mlp": |
| add(prefix + ".mlp_alpha", (h,)) |
| for name in ("gate_proj", "up_proj"): |
| add(prefix + f".mixer.{name}.weight", (f, h)) |
| add(prefix + ".mixer.down_proj.weight", (h, f)) |
| for name in ("suv_gate", "suv_up"): |
| add(prefix + f".mixer.{name}", (f,)) |
| else: |
| add(prefix + ".attn_alpha", (h,)) |
| for name, rows in (("q_proj", q), ("k_proj", kv), ("v_proj", kv), ("g_proj", q)): |
| add(prefix + f".mixer.{name}.weight", (rows, h)) |
| add(prefix + ".mixer.o_proj.weight", (h, q)) |
| for name in ("sqk", "s_gate"): |
| add(prefix + f".mixer.{name}", (q,)) |
| save_file(tensors, str(folder / "model.safetensors"), metadata={"format": "pt"}) |
| return str(graph) |
| |
| |
| create_model( |
| source_model_id="nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-BF16", |
| source_revision="b366b87f762ea334ba3eb55912031baba1515810", |
| save_folder="/tmp/dory-1.0", |
| ) |
| ``` |
|
|
| </details> |
|
|
| ### Printing the model: |
|
|
| <details><summary>Click to expand</summary> |
|
|
| ```text |
| CheckpointModule( |
| (backbone): CheckpointModule( |
| (embeddings): CheckpointModule(weight: shape=(131072, 2), dtype=torch.bfloat16) |
| (layers): CheckpointModule( |
| (0): CheckpointModule( |
| attn_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| sqk: shape=(128,), dtype=torch.bfloat16 |
| s_gate: shape=(128,), dtype=torch.bfloat16 |
| (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16) |
| ) |
| ) |
| (1): CheckpointModule( |
| mlp_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| suv_gate: shape=(32,), dtype=torch.bfloat16 |
| suv_up: shape=(32,), dtype=torch.bfloat16 |
| (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16) |
| ) |
| ) |
| (2): CheckpointModule( |
| attn_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| sqk: shape=(128,), dtype=torch.bfloat16 |
| s_gate: shape=(128,), dtype=torch.bfloat16 |
| (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16) |
| ) |
| ) |
| (3): CheckpointModule( |
| mlp_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| suv_gate: shape=(32,), dtype=torch.bfloat16 |
| suv_up: shape=(32,), dtype=torch.bfloat16 |
| (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16) |
| ) |
| ) |
| (4): CheckpointModule( |
| attn_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| sqk: shape=(128,), dtype=torch.bfloat16 |
| s_gate: shape=(128,), dtype=torch.bfloat16 |
| (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16) |
| ) |
| ) |
| (5): CheckpointModule( |
| mlp_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| suv_gate: shape=(32,), dtype=torch.bfloat16 |
| suv_up: shape=(32,), dtype=torch.bfloat16 |
| (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16) |
| ) |
| ) |
| (6): CheckpointModule( |
| attn_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| sqk: shape=(128,), dtype=torch.bfloat16 |
| s_gate: shape=(128,), dtype=torch.bfloat16 |
| (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16) |
| ) |
| ) |
| (7): CheckpointModule( |
| mlp_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| suv_gate: shape=(32,), dtype=torch.bfloat16 |
| suv_up: shape=(32,), dtype=torch.bfloat16 |
| (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16) |
| ) |
| ) |
| (8): CheckpointModule( |
| attn_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| sqk: shape=(128,), dtype=torch.bfloat16 |
| s_gate: shape=(128,), dtype=torch.bfloat16 |
| (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16) |
| ) |
| ) |
| (9): CheckpointModule( |
| mlp_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| suv_gate: shape=(32,), dtype=torch.bfloat16 |
| suv_up: shape=(32,), dtype=torch.bfloat16 |
| (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16) |
| ) |
| ) |
| (10): CheckpointModule( |
| attn_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| sqk: shape=(128,), dtype=torch.bfloat16 |
| s_gate: shape=(128,), dtype=torch.bfloat16 |
| (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16) |
| ) |
| ) |
| (11): CheckpointModule( |
| mlp_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| suv_gate: shape=(32,), dtype=torch.bfloat16 |
| suv_up: shape=(32,), dtype=torch.bfloat16 |
| (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16) |
| ) |
| ) |
| (12): CheckpointModule( |
| attn_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| sqk: shape=(128,), dtype=torch.bfloat16 |
| s_gate: shape=(128,), dtype=torch.bfloat16 |
| (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16) |
| ) |
| ) |
| (13): CheckpointModule( |
| mlp_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| suv_gate: shape=(32,), dtype=torch.bfloat16 |
| suv_up: shape=(32,), dtype=torch.bfloat16 |
| (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16) |
| ) |
| ) |
| (14): CheckpointModule( |
| attn_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| sqk: shape=(128,), dtype=torch.bfloat16 |
| s_gate: shape=(128,), dtype=torch.bfloat16 |
| (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16) |
| ) |
| ) |
| (15): CheckpointModule( |
| mlp_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| suv_gate: shape=(32,), dtype=torch.bfloat16 |
| suv_up: shape=(32,), dtype=torch.bfloat16 |
| (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16) |
| ) |
| ) |
| (16): CheckpointModule( |
| attn_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| sqk: shape=(128,), dtype=torch.bfloat16 |
| s_gate: shape=(128,), dtype=torch.bfloat16 |
| (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16) |
| ) |
| ) |
| (17): CheckpointModule( |
| mlp_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| suv_gate: shape=(32,), dtype=torch.bfloat16 |
| suv_up: shape=(32,), dtype=torch.bfloat16 |
| (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16) |
| ) |
| ) |
| (18): CheckpointModule( |
| attn_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| sqk: shape=(128,), dtype=torch.bfloat16 |
| s_gate: shape=(128,), dtype=torch.bfloat16 |
| (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16) |
| ) |
| ) |
| (19): CheckpointModule( |
| mlp_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| suv_gate: shape=(32,), dtype=torch.bfloat16 |
| suv_up: shape=(32,), dtype=torch.bfloat16 |
| (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16) |
| ) |
| ) |
| (20): CheckpointModule( |
| attn_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| sqk: shape=(128,), dtype=torch.bfloat16 |
| s_gate: shape=(128,), dtype=torch.bfloat16 |
| (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16) |
| ) |
| ) |
| (21): CheckpointModule( |
| mlp_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| suv_gate: shape=(32,), dtype=torch.bfloat16 |
| suv_up: shape=(32,), dtype=torch.bfloat16 |
| (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16) |
| ) |
| ) |
| (22): CheckpointModule( |
| attn_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| sqk: shape=(128,), dtype=torch.bfloat16 |
| s_gate: shape=(128,), dtype=torch.bfloat16 |
| (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16) |
| ) |
| ) |
| (23): CheckpointModule( |
| mlp_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| suv_gate: shape=(32,), dtype=torch.bfloat16 |
| suv_up: shape=(32,), dtype=torch.bfloat16 |
| (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16) |
| ) |
| ) |
| (24): CheckpointModule( |
| attn_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| sqk: shape=(128,), dtype=torch.bfloat16 |
| s_gate: shape=(128,), dtype=torch.bfloat16 |
| (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16) |
| ) |
| ) |
| (25): CheckpointModule( |
| mlp_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| suv_gate: shape=(32,), dtype=torch.bfloat16 |
| suv_up: shape=(32,), dtype=torch.bfloat16 |
| (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16) |
| ) |
| ) |
| (26): CheckpointModule( |
| attn_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| sqk: shape=(128,), dtype=torch.bfloat16 |
| s_gate: shape=(128,), dtype=torch.bfloat16 |
| (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16) |
| ) |
| ) |
| (27): CheckpointModule( |
| mlp_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| suv_gate: shape=(32,), dtype=torch.bfloat16 |
| suv_up: shape=(32,), dtype=torch.bfloat16 |
| (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16) |
| ) |
| ) |
| (28): CheckpointModule( |
| attn_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| sqk: shape=(128,), dtype=torch.bfloat16 |
| s_gate: shape=(128,), dtype=torch.bfloat16 |
| (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16) |
| ) |
| ) |
| (29): CheckpointModule( |
| mlp_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| suv_gate: shape=(32,), dtype=torch.bfloat16 |
| suv_up: shape=(32,), dtype=torch.bfloat16 |
| (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16) |
| ) |
| ) |
| (30): CheckpointModule( |
| attn_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| sqk: shape=(128,), dtype=torch.bfloat16 |
| s_gate: shape=(128,), dtype=torch.bfloat16 |
| (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16) |
| ) |
| ) |
| (31): CheckpointModule( |
| mlp_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| suv_gate: shape=(32,), dtype=torch.bfloat16 |
| suv_up: shape=(32,), dtype=torch.bfloat16 |
| (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16) |
| ) |
| ) |
| (32): CheckpointModule( |
| attn_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| sqk: shape=(128,), dtype=torch.bfloat16 |
| s_gate: shape=(128,), dtype=torch.bfloat16 |
| (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16) |
| ) |
| ) |
| (33): CheckpointModule( |
| mlp_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| suv_gate: shape=(32,), dtype=torch.bfloat16 |
| suv_up: shape=(32,), dtype=torch.bfloat16 |
| (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16) |
| ) |
| ) |
| (34): CheckpointModule( |
| attn_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| sqk: shape=(128,), dtype=torch.bfloat16 |
| s_gate: shape=(128,), dtype=torch.bfloat16 |
| (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16) |
| ) |
| ) |
| (35): CheckpointModule( |
| mlp_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| suv_gate: shape=(32,), dtype=torch.bfloat16 |
| suv_up: shape=(32,), dtype=torch.bfloat16 |
| (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16) |
| ) |
| ) |
| (36): CheckpointModule( |
| attn_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| sqk: shape=(128,), dtype=torch.bfloat16 |
| s_gate: shape=(128,), dtype=torch.bfloat16 |
| (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16) |
| ) |
| ) |
| (37): CheckpointModule( |
| mlp_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| suv_gate: shape=(32,), dtype=torch.bfloat16 |
| suv_up: shape=(32,), dtype=torch.bfloat16 |
| (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16) |
| ) |
| ) |
| (38): CheckpointModule( |
| attn_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| sqk: shape=(128,), dtype=torch.bfloat16 |
| s_gate: shape=(128,), dtype=torch.bfloat16 |
| (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16) |
| ) |
| ) |
| (39): CheckpointModule( |
| mlp_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| suv_gate: shape=(32,), dtype=torch.bfloat16 |
| suv_up: shape=(32,), dtype=torch.bfloat16 |
| (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16) |
| ) |
| ) |
| (40): CheckpointModule( |
| attn_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| sqk: shape=(128,), dtype=torch.bfloat16 |
| s_gate: shape=(128,), dtype=torch.bfloat16 |
| (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16) |
| ) |
| ) |
| (41): CheckpointModule( |
| mlp_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| suv_gate: shape=(32,), dtype=torch.bfloat16 |
| suv_up: shape=(32,), dtype=torch.bfloat16 |
| (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16) |
| ) |
| ) |
| (42): CheckpointModule( |
| attn_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| sqk: shape=(128,), dtype=torch.bfloat16 |
| s_gate: shape=(128,), dtype=torch.bfloat16 |
| (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16) |
| ) |
| ) |
| (43): CheckpointModule( |
| mlp_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| suv_gate: shape=(32,), dtype=torch.bfloat16 |
| suv_up: shape=(32,), dtype=torch.bfloat16 |
| (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16) |
| ) |
| ) |
| (44): CheckpointModule( |
| attn_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| sqk: shape=(128,), dtype=torch.bfloat16 |
| s_gate: shape=(128,), dtype=torch.bfloat16 |
| (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16) |
| ) |
| ) |
| (45): CheckpointModule( |
| mlp_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| suv_gate: shape=(32,), dtype=torch.bfloat16 |
| suv_up: shape=(32,), dtype=torch.bfloat16 |
| (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16) |
| ) |
| ) |
| (46): CheckpointModule( |
| attn_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| sqk: shape=(128,), dtype=torch.bfloat16 |
| s_gate: shape=(128,), dtype=torch.bfloat16 |
| (q_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (k_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (v_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (g_proj): CheckpointModule(weight: shape=(128, 2), dtype=torch.bfloat16) |
| (o_proj): CheckpointModule(weight: shape=(2, 128), dtype=torch.bfloat16) |
| ) |
| ) |
| (47): CheckpointModule( |
| mlp_alpha: shape=(2,), dtype=torch.bfloat16 |
| (mixer): CheckpointModule( |
| suv_gate: shape=(32,), dtype=torch.bfloat16 |
| suv_up: shape=(32,), dtype=torch.bfloat16 |
| (gate_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (up_proj): CheckpointModule(weight: shape=(32, 2), dtype=torch.bfloat16) |
| (down_proj): CheckpointModule(weight: shape=(2, 32), dtype=torch.bfloat16) |
| ) |
| ) |
| ) |
| ) |
| (lm_head): CheckpointModule( |
| weight: shape=(131072, 2), dtype=torch.bfloat16 |
| logit_scale: shape=(131072,), dtype=torch.bfloat16 |
| ) |
| ) |
| ``` |
|
|
| </details> |
|
|
| Environment: [.meta.json](.meta.json) |
|
|
| ### Change log: |
|
|
| - 2026-10-25: Created BF16 Dory with a two-dimensional latent and verified TP=4 inference. |
| - 2026-10-26: Verified the current standalone offline and HTTP serving examples on four H20 GPUs. |
|
|
| This repo is generated by GPT-6.1-Sol |
