AIhomeJP/home1

TRM-GPT

TRM-GPT is a compact causal language model implemented with PyTorch and Hugging Face Transformers.

The model combines standard Transformer blocks with a recursive refinement block. It is loaded as a custom AutoModelForCausalLM implementation using trust_remote_code=True.

Model architecture

The model contains:

• Token embeddings
• Learned positional embeddings
• Causal self-attention blocks
• Feed-forward MLP blocks
• Recursive refinement block
• Final layer normalization
• Linear language-modeling head

The output projection is available as:

model.lm_head

Loading the model

Because this repository contains custom model code, trust_remote_code=True is required.

from transformers import AutoModelForCausalLM, AutoTokenizer

repo_id = "AIhomeJP/home1"

tokenizer = AutoTokenizer.from_pretrained(
    repo_id,
    trust_remote_code=True,
)

model = AutoModelForCausalLM.from_pretrained(
    repo_id,
    trust_remote_code=True,
)

print(type(model))
print(model.lm_head)

Text generation

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo_id = "AIhomeJP/home1"

tokenizer = AutoTokenizer.from_pretrained(
    repo_id,
    trust_remote_code=True,
)

model = AutoModelForCausalLM.from_pretrained(
    repo_id,
    trust_remote_code=True,
)

device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
model.eval()

inputs = tokenizer(
    "こんにちは",
    return_tensors="pt",
).to(device)

with torch.no_grad():
    output_ids = model.generate(
        **inputs,
        max_new_tokens=50,
        do_sample=True,
        temperature=0.8,
        top_k=50,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.eos_token_id,
    )

print(
    tokenizer.decode(
        output_ids[0],
        skip_special_tokens=True,
    )
)

About lm_head.weight

This model uses tied input and output embeddings.

The following two parameters share the same underlying weight:

token_emb.weight
lm_head.weight

The model class must declare the shared parameter:

class TRMGPTForCausalLM(PreTrainedModel):
    config_class = TRMGPTConfig
    base_model_prefix = "trm_gpt"

    _tied_weights_keys = ["lm_head.weight"]

The model also exposes the output embedding through the Transformers interface:

def get_output_embeddings(self):
    return self.lm_head

def set_output_embeddings(self, new_embeddings):
    self.lm_head = new_embeddings

Weight tying is performed with:

def tie_weights(self):
    super().tie_weights()

    if self.config.tie_word_embeddings:
        self._tie_or_clone_weights(
            self.lm_head,
            self.token_emb,
        )

Why lm_head.weight may not appear in the Safetensors file

When token_emb.weight and lm_head.weight share the same storage, Safetensors may save only one copy of the tensor.

Therefore, the absence of a separate lm_head.weight entry in model.safetensors does not necessarily mean that the model has no lm_head.

After loading, Transformers restores the shared relationship between:

token_emb.weight
lm_head.weight

The relationship can be checked with:

assert hasattr(model, "lm_head")

assert (
    model.lm_head.weight.data_ptr()
    == model.token_emb.weight.data_ptr()
)

print("lm_head is available and correctly tied")

Troubleshooting

lm_head.weight is reported as missing

Update the custom model class so that it contains:

_tied_weights_keys = ["lm_head.weight"]

For compatibility with some Transformers versions, the following declaration can also be added:

_keys_to_ignore_on_load_missing = [
    r"lm_head\.weight",
]

Then load and save the model again:

from transformers import AutoModelForCausalLM, AutoTokenizer

repo_id = "AIhomeJP/home1"
output_dir = "./home1_fixed"

model = AutoModelForCausalLM.from_pretrained(
    repo_id,
    trust_remote_code=True,
)

tokenizer = AutoTokenizer.from_pretrained(
    repo_id,
    trust_remote_code=True,
)

model.tie_weights()

assert (
    model.lm_head.weight.data_ptr()
    == model.token_emb.weight.data_ptr()
)

model.save_pretrained(
    output_dir,
    safe_serialization=True,
)

tokenizer.save_pretrained(output_dir)

Custom model code is not executed

Make sure that trust_remote_code=True is specified:

model = AutoModelForCausalLM.from_pretrained(
    "AIhomeJP/home1",
    trust_remote_code=True,
)

Sequence length exceeds the model limit

Input length must not exceed config.block_size.

inputs = tokenizer(
    text,
    truncation=True,
    max_length=model.config.block_size,
    return_tensors="pt",
)

Tokenizer has no padding token

For generation, the EOS token can be used as the padding token:

if tokenizer.pad_token_id is None:
    tokenizer.pad_token = tokenizer.eos_token

Training

For causal language modeling, labels are normally the same token sequence as input_ids.

outputs = model(
    input_ids=input_ids,
    attention_mask=attention_mask,
    labels=input_ids,
)

loss = outputs.loss
loss.backward()

The model implementation shifts logits and labels internally:

shift_logits = logits[:, :-1, :].contiguous()
shift_labels = labels[:, 1:].contiguous()

Padding or ignored target positions should use label value -100.

Configuration

Example configuration:

{
  "model_type": "trm_gpt",
  "vocab_size": 50257,
  "block_size": 256,
  "n_layer": 4,
  "n_embd": 256,
  "n_head": 4,
  "dropout": 0.1,
  "recursive_steps": 2,
  "tie_word_embeddings": true
}

Security

This repository uses custom Python code.

Review the repository files before enabling:

trust_remote_code=True

License

MIT License