A transformer-based language model with hierarchical memory architecture for unlimited context.
Redformer implements a novel multi-level memory system that enables processing arbitrarily long sequences:
- Hierarchical Memory: 15 levels of memory (2^0, 2^1, ..., 2^14 tokens) for efficient long-range context
- TwoPacker: Geometric algebra-based operations for memory compression
- Reducer: Attention-based memory reduction across levels
- Memory Persistence: Memory state carries across sequences for continuous learning
pip install -r requirements.txtTrain the model using trainer2.py:
python trainer2.pyThe trainer uses Hugging Face Accelerate for distributed training with mixed precision (bfloat16). Models are saved to ~/ReD2B/ by default.
Training Details:
- Memory updates happen after optimization step (not in forward pass) to prevent double-updates with gradient checkpointing
- Supports streaming datasets for large-scale pretraining
- Automatic checkpointing every 1000 steps
Run the interactive chat interface:
python demo_chat.pyFeatures:
- Streaming token generation
- Loss display showing model "surprise"
- Multi-line input support (Alt-Enter or Ctrl-J)
- Persistent memory across conversation turns
Manage model checkpoints with Google Cloud Storage:
# list
./rocli list
# Upload local model to GCS
./rocli upload ~/ReD2B/
# List stored models
./rocli list
# Download model from GCS
./rocli download ReD2B ~/
# Delete stored model
./rocli delete Red2BThe CLI provides a lightweight interface for syncing large model files to cloud storage without external dependencies.
modelling/model.py- Core model and memory wrappermodelling/layer.py- Layer architecture with hierarchical memorymodelling/twopack.py- Geometric packing operationsmodelling/reducer.py- Attention-based reductionmodelling/config.py- Model configuration
Default config (2.8B parameters + memory):
hidden_size = 2816
num_packer_heads = 256
num_reducer_heads = 128
num_layers = 8
memory_levels = 15