Upload 7 files

comparison/config_mamba.json

@@ -0,0 +1,38 @@
+{
+  "architectures": [
+    "MambaForCausalLM"
+  ],
+  "bos_token_id": 0,
+  "conv_kernel": 4,
+  "d_inner": 1024,
+  "d_model": 512,
+  "eos_token_id": 0,
+  "expand": 2,
+  "fused_add_norm": true,
+  "hidden_act": "silu",
+  "hidden_size": 512,
+  "initializer_range": 0.1,
+  "intermediate_size": 1024,
+  "layer_norm_epsilon": 1e-05,
+  "model_type": "mamba",
+  "n_layer": 18,
+  "num_hidden_layers": 18,
+  "pad_token_id": 0,
+  "pad_vocab_size_multiple": 8,
+  "rescale_prenorm_residual": false,
+  "residual_in_fp32": true,
+  "rms_norm": true,
+  "ssm_cfg": {},
+  "state_size": 16,
+  "time_step_floor": 0.0001,
+  "time_step_init_scheme": "random",
+  "time_step_max": 0.1,
+  "time_step_min": 0.001,
+  "time_step_rank": 32,
+  "time_step_scale": 1.0,
+  "torch_dtype": "bfloat16",
+  "use_bias": false,
+  "use_cache": true,
+  "use_conv_bias": true,
+  "vocab_size": 264
+}

comparison/config_pmnet.json

@@ -0,0 +1,58 @@
+{
+  "architectures": [
+    "PMNetForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.1,
+  "bos_token_id": 0,
+  "dtype": "bfloat16",
+  "eos_token_id": 0,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 384,
+  "initializer_range": 0.02,
+  "intermediate_size": 1024,
+  "layer_types": [
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention"
+  ],
+  "max_position_embeddings": 40960,
+  "max_window_layers": 0,
+  "memory_cumsum": true,
+  "memory_size": 32,
+  "memory_write_period": 3,
+  "model_type": "pmnet",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "num_key_value_heads": 12,
+  "num_memory": 4,
+  "num_memory_read_heads": 4,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000,
+  "sliding_window": 128,
+  "tie_word_embeddings": true,
+  "transformers_version": "4.57.6",
+  "use_cache": false,
+  "use_sliding_window": true,
+  "vocab_size": 264,
+  "auto_map": {
+    "AutoConfig": "configuration_pmnet.PMNetConfig",
+    "AutoModel": "modeling_pmnet.PMNetModel",
+    "AutoModelForCausalLM": "modeling_pmnet.PMNetForCausalLM",
+    "AutoTokenizer": [
+      "tokenization_pmnet.ByteTokenizer",
+      null
+    ]
+  }
+}

comparison/config_pmnet_no_mem.json

@@ -0,0 +1,58 @@
+{
+  "architectures": [
+    "PMNetForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.1,
+  "bos_token_id": 0,
+  "dtype": "bfloat16",
+  "eos_token_id": 0,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 384,
+  "initializer_range": 0.02,
+  "intermediate_size": 1024,
+  "layer_types": [
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention"
+  ],
+  "max_position_embeddings": 40960,
+  "max_window_layers": 0,
+  "memory_cumsum": false,
+  "memory_size": 32,
+  "memory_write_period": 3,
+  "model_type": "pmnet",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "num_key_value_heads": 12,
+  "num_memory": 4,
+  "num_memory_read_heads": 4,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000,
+  "sliding_window": 128,
+  "tie_word_embeddings": true,
+  "transformers_version": "4.57.6",
+  "use_cache": false,
+  "use_sliding_window": true,
+  "vocab_size": 264,
+  "auto_map": {
+    "AutoConfig": "configuration_pmnet.PMNetConfig",
+    "AutoModel": "modeling_pmnet.PMNetModel",
+    "AutoModelForCausalLM": "modeling_pmnet.PMNetForCausalLM",
+    "AutoTokenizer": [
+      "tokenization_pmnet.ByteTokenizer",
+      null
+    ]
+  }
+}

comparison/config_smollm.json

@@ -0,0 +1,29 @@
+{
+  "_name_or_path": "llama-30M-config",
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.1,
+  "bos_token_id": 0,
+  "eos_token_id": 0,
+  "hidden_act": "silu",
+  "hidden_size": 384,
+  "initializer_range": 0.02,
+  "intermediate_size": 1024,
+  "max_position_embeddings": 8192,
+  "mlp_bias": false,
+  "model_type": "llama",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 16,
+  "num_key_value_heads": 4,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_scaling": null,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.41.2",
+  "use_cache": true,
+  "vocab_size": 264
+}

comparison/fineweb-edu_1B_2048.py

@@ -0,0 +1,95 @@
+
+import wandb
+import os
+os.environ["HF_HOME"] = "/root/hf_cache"
+os.environ["HF_DATASETS_CACHE"] = "/root/hf_cache/datasets"
+from transformers import AutoTokenizer
+NUM_PROC = 16
+BLOCK_SIZE = 2048
+DATA_DIR = "/data/fineweb_edu_byte_1B_2048"
+SHARD_START_INDEX = 0
+SHARD_LEN = 1
+DATA_PROCESS_BATCH_SIZE=5000
+dataset_base_url = "https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu/resolve/main/sample/10BT"
+
+from datasets import load_dataset, load_from_disk
+from itertools import chain
+
+
+from typing import Dict, List, Optional, Any
+
+tokenizer = AutoTokenizer.from_pretrained("google/byt5-small")
+
+train_files = {
+    "train": [
+        f"{dataset_base_url}/{i:03d}_00000.parquet" for i in range(SHARD_START_INDEX, SHARD_START_INDEX+SHARD_LEN)
+    ]
+}
+raw_train = load_dataset("parquet", data_files=train_files, split="train") # FIXME
+
+val_files = {"train": f"{dataset_base_url}/013_00000.parquet"}
+raw_val = load_dataset("parquet", data_files=val_files, split="train") # FIXME
+
+def process_batch(examples):
+    tokenized = tokenizer(examples["text"], add_special_tokens=False)
+    input_ids = tokenized["input_ids"]
+
+    eos_id = tokenizer.eos_token_id
+    concatenated_ids = list(chain.from_iterable(ids + [eos_id] for ids in input_ids))
+
+    total_length = len(concatenated_ids)
+    total_length = (total_length // BLOCK_SIZE) * BLOCK_SIZE
+
+    result = {
+        "input_ids": [
+            concatenated_ids[i : i + BLOCK_SIZE]
+            for i in range(0, total_length, BLOCK_SIZE)
+        ]
+    }
+    return result
+
+column_names = raw_train.column_names
+
+print(f"Processing Train dataset using {NUM_PROC} processes...")
+lm_dataset_train = raw_train.map(
+    process_batch,
+    batched=True,
+    batch_size=DATA_PROCESS_BATCH_SIZE,
+    remove_columns=column_names,
+    num_proc=NUM_PROC,
+    desc="Tokenizing & Grouping Train"
+)
+
+print("Processing Val dataset...")
+lm_dataset_val = raw_val.map(
+    process_batch,
+    batched=True,
+    batch_size=DATA_PROCESS_BATCH_SIZE,
+    remove_columns=column_names,
+    num_proc=NUM_PROC,
+    desc="Tokenizing & Grouping Val"
+)
+
+os.makedirs(DATA_DIR, exist_ok=True)
+
+print(f"Saving to {DATA_DIR}...")
+TARGET_BLOCKS_TRAIN = 250_000
+TARGET_BLOCKS_VAL = 1_250
+TARGET_BLOCKS_TEST = 50_000
+
+lm_dataset_train = lm_dataset_train.select(range(TARGET_BLOCKS_TRAIN))
+lm_dataset_test = lm_dataset_val.select(range(TARGET_BLOCKS_TEST))
+lm_dataset_val = lm_dataset_val.select(range(TARGET_BLOCKS_TEST, TARGET_BLOCKS_TEST + TARGET_BLOCKS_VAL))
+
+print(f"Saving finalized datasets to {DATA_DIR}...")
+os.makedirs(DATA_DIR, exist_ok=True)
+lm_dataset_train.save_to_disk(f"{DATA_DIR}/train")
+lm_dataset_val.save_to_disk(f"{DATA_DIR}/val")
+lm_dataset_test.save_to_disk(f"{DATA_DIR}/test")
+
+ds_train = load_from_disk(f"{DATA_DIR}/train")
+print(f"Final Train blocks count: {len(ds_train)}")
+ds_val = load_from_disk(f"{DATA_DIR}/val")
+print(f"Final Val blocks count: {len(ds_val)}")
+ds_test = load_from_disk(f"{DATA_DIR}/test")
+print(f"Final Test blocks count: {len(ds_test)}")

comparison/install_dependency.sh

@@ -0,0 +1,2 @@
+pip install transformers[torch]==4.57.6
+pip install wandb datasets einops hf_transfer accelerate https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.8cxx11abiTRUE-cp312-cp312-linux_x86_64.whl https://github.com/state-spaces/mamba/releases/download/v2.3.1/mamba_ssm-2.3.1+cu12torch2.8cxx11abiTRUE-cp312-cp312-linux_x86_64.whl https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.6.1.post4/causal_conv1d-1.6.1+cu12torch2.8cxx11abiTRUE-cp312-cp312-linux_x86_64.whl

comparison/train_comparison.py

@@ -0,0 +1,143 @@
+import os
+import json
+import torch
+import wandb
+from datasets import load_from_disk
+from transformers import (
+    Trainer,
+    TrainingArguments,
+    DataCollatorForLanguageModeling,
+    set_seed,
+    get_cosine_schedule_with_warmup,
+    AutoModelForCausalLM, AutoTokenizer,AutoConfig
+)
+from torch.optim import AdamW
+
+# MODEL_ID = "pmnet-icml2026-10724/pmnet"
+# CONFIG_PATH = "/workspace/scripts/config_pmnet.json"
+# CKPT_SAVE_DIR = "/workspace/ckpts/pmnet"
+# compile_model = True
+
+# MODEL_ID = "state-spaces/mamba-130m-hf"
+# CONFIG_PATH = "/workspace/scripts/config_mamba.json"
+# CKPT_SAVE_DIR = "/workspace/ckpts/mamba"
+# compile_model = False
+
+# MODEL_ID = "HuggingFaceTB/SmolLM-135M"
+# CONFIG_PATH = "/workspace/scripts/config_smollm.json"
+# CKPT_SAVE_DIR = "/workspace/ckpts/smollm"
+# compile_model = True
+
+MODEL_ID = "pmnet-icml2026-10724/pmnet"
+CONFIG_PATH = "/workspace/scripts/config_pmnet_no_mem.json"
+CKPT_SAVE_DIR = "/workspace/ckpts/pmnet_no_mem"
+compile_model = True
+
+########################################################################
+
+DATA_DIR = "/data/fineweb_edu_byte_1B_2048"
+SEED = 42
+
+BATCH_SIZE = 48
+GRADIENT_ACCUMULATION_STEPS = 1
+NUM_DEVICES = 2
+
+os.environ["WANDB_PROJECT"] = "pmnet_comparison"
+
+########################################################################
+
+per_device_batch_size = BATCH_SIZE // (NUM_DEVICES*GRADIENT_ACCUMULATION_STEPS)
+
+
+def get_optimizer_grouped_parameters(model, weight_decay: float):
+    no_decay_keywords = ["bias", "norm", "embedding", "layernorm", "a_log"]
+    decay_params = []
+    no_decay_params = []
+
+    for name, param in model.named_parameters():
+        if not param.requires_grad: continue
+        if any(keyword in name.lower() for keyword in no_decay_keywords):
+            no_decay_params.append(param)
+        else:
+            decay_params.append(param)
+    return [
+        {"params": decay_params, "weight_decay": weight_decay},
+        {"params": no_decay_params, "weight_decay": 0.0},
+    ]
+
+def main():
+    set_seed(SEED)
+
+    
+    config = AutoConfig.from_pretrained(MODEL_ID, trust_remote_code=True)
+    if CONFIG_PATH and os.path.exists(CONFIG_PATH):
+        with open(CONFIG_PATH, "r", encoding="utf-8") as f:
+            local_config_dict = json.load(f)
+        config.update(local_config_dict)
+    model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
+
+    # model size
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Total params: {total_params/1_000_000}")
+
+    dataset_train = load_from_disk(os.path.join(DATA_DIR, "train"))
+    dataset_val = load_from_disk(os.path.join(DATA_DIR, "val"))
+
+    tokenizer = AutoTokenizer.from_pretrained("google/byt5-small")
+
+    training_args = TrainingArguments(
+        output_dir=CKPT_SAVE_DIR,
+        overwrite_output_dir=False,
+        num_train_epochs=1,
+
+        save_strategy="steps",
+        save_steps=500,
+        eval_strategy="steps",
+        eval_steps=250,
+        logging_strategy="steps",
+        logging_steps=10,
+        save_total_limit=2,
+
+        load_best_model_at_end=True,
+        metric_for_best_model="loss",
+        greater_is_better=False,
+
+        seed=SEED,
+        data_seed=SEED,
+
+        per_device_train_batch_size=per_device_batch_size,
+        per_device_eval_batch_size=per_device_batch_size,
+        gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS,
+
+        learning_rate=1e-4,
+        max_grad_norm=1.0,
+        fp16=False,
+        bf16=True,
+
+        dataloader_num_workers=8,
+        report_to="wandb",
+        ddp_find_unused_parameters=False,
+
+        lr_scheduler_type="cosine",
+        warmup_steps=500,
+        torch_compile=compile_model
+    )
+
+    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+    optimizer_grouped_parameters = get_optimizer_grouped_parameters(model, weight_decay=0.1)
+    optimizer = AdamW(optimizer_grouped_parameters, lr=training_args.learning_rate, betas=(0.9, 0.95))
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=dataset_train,
+        eval_dataset=dataset_val,
+        data_collator=data_collator,
+        optimizers=(optimizer, None),
+    )
+    trainer.train()
+    trainer.save_model(os.path.join(CKPT_SAVE_DIR, "final_model"))
+
+if __name__ == "__main__":
+    main()