Pre-compiled and pre-sharded
Trinity-Nano-Preview (~6B total, ~1B active MoE) for AWS Neuron SDK 2.28, ready to load on
inf2.xlarge (16GB system RAM) or any larger Inferentia2/Trainium instance.
The standard NxDI load path downloads the full HuggingFace checkpoint (~12GB bf16) into CPU RAM for weight conversion and sharding. On inf2.xlarge (16GB system RAM), this causes an OOM kill at 15+ GB RSS.
Pre-sharded weights bypass this entirely — NxDI reads directly from the per-rank sharded files, using only 1.4 GB RSS (12.6% of system RAM).
The Trinity Neuron implementation is not yet merged into the main NxDI repo. Use the contrib branch from the fork:
1git clone --branch contrib/trinity-model --single-branch \
2 https://github.com/jimburtoft/neuronx-distributed-inference.git nxdi-trinity
1from huggingface_hub import snapshot_download
2
3# Download the pre-compiled artifact (model.pt + sharded weights)
4snapshot_download("jburtoft/Trinity-Nano-Neuron-TP1",
5 local_dir="/home/ubuntu/Trinity-Nano-Neuron-TP1")
6
7# Download config + tokenizer only (no model weights needed)
8snapshot_download("arcee-ai/Trinity-Nano-Preview",
9 local_dir="/home/ubuntu/Trinity-Nano-Preview",
10 ignore_patterns=["*.safetensors", "*.bin", "*.pt", "*.gguf"])
1import sys
2import torch
3from transformers import AutoTokenizer
4from neuronx_distributed_inference.models.config import MoENeuronConfig
5
6# Point to the Trinity implementation from the cloned repo
7sys.path.insert(0, "/home/ubuntu/nxdi-trinity/contrib/models/Trinity/src")
8from modeling_trinity import NeuronTrinityForCausalLM, TrinityInferenceConfig
9
10# Build model with save_sharded_checkpoint=True (must match compilation)
11neuron_config = MoENeuronConfig(
12 tp_degree=1,
13 batch_size=1,
14 seq_len=2048,
15 torch_dtype=torch.bfloat16,
16 save_sharded_checkpoint=True,
17)
18
19config = TrinityInferenceConfig.from_pretrained(
20 "/home/ubuntu/Trinity-Nano-Preview",
21 neuron_config=neuron_config,
22)
23
24model = NeuronTrinityForCausalLM("/home/ubuntu/Trinity-Nano-Preview", config)
25model.load("/home/ubuntu/Trinity-Nano-Neuron-TP1")
26
27# Tokenize
28tokenizer = AutoTokenizer.from_pretrained(
29 "/home/ubuntu/Trinity-Nano-Preview", trust_remote_code=True
30)
31
32prompt = "Hello, how are you today?"
33inputs = tokenizer(prompt, return_tensors="pt")
34input_ids = inputs.input_ids
35
36# Generate
37model.reset()
38position_ids = torch.arange(input_ids.shape[1]).unsqueeze(0)
39seq_ids = torch.arange(1)
40
41with torch.no_grad():
42 outputs = model(input_ids, position_ids=position_ids, seq_ids=seq_ids)
43
44logits = outputs.logits if hasattr(outputs, "logits") else outputs[0]
45next_token = torch.argmax(logits[:, -1, :], dim=-1)
46print(f"Prompt: {prompt}")
47print(f"Next token: {tokenizer.decode(next_token)}")
48
49# Autoregressive generation
50generated = [next_token.unsqueeze(0)]
51for i in range(31):
52 pos = torch.tensor([[input_ids.shape[1] + i]])
53 with torch.no_grad():
54 outputs = model(generated[-1], position_ids=pos, seq_ids=seq_ids)
55 logits = outputs.logits if hasattr(outputs, "logits") else outputs[0]
56 next_token = torch.argmax(logits[:, -1, :], dim=-1)
57 generated.append(next_token.unsqueeze(0))
58
59text = tokenizer.decode(torch.cat(generated, dim=1)[0], skip_special_tokens=True)
60print(f"Generated: {text}")
To compile for different configurations (e.g., TP=2, BS=4), you need a larger instance (inf2.8xlarge or trn2.3xlarge):
1import sys
2import torch
3from neuronx_distributed_inference.models.config import MoENeuronConfig
4
5sys.path.insert(0, "/path/to/nxdi-trinity/contrib/models/Trinity/src")
6from modeling_trinity import NeuronTrinityForCausalLM, TrinityInferenceConfig
7
8neuron_config = MoENeuronConfig(
9 tp_degree=1, # Adjust as needed
10 batch_size=1, # Adjust as needed
11 seq_len=2048, # Adjust as needed
12 torch_dtype=torch.bfloat16,
13 save_sharded_checkpoint=True, # Required for pre-sharded deployment
14)
15
16config = TrinityInferenceConfig.from_pretrained(
17 "/path/to/Trinity-Nano-Preview", neuron_config=neuron_config
18)
19model = NeuronTrinityForCausalLM("/path/to/Trinity-Nano-Preview", config)
20model.compile("/path/to/compiled-output")
21# Output: model.pt, neuron_config.json, weights/tp{rank}_sharded_checkpoint.safetensors
The NeuronX Distributed Inference implementation for Trinity is available at:
github.com/jimburtoft/neuronx-distributed-inference (branch:
contrib/trinity-model)
This implementation supports all three Trinity model sizes (Nano, Mini, Large) with a single unified modeling_trinity.py.
