ARMT-Gemma-3-1B-IT-8x1024 (portable copy)

Same weights as irodkin/ARMT-Gemma-3-1B-IT-8x1024 (Associative Recurrent Memory Transformer on top of Gemma-3-1B-IT, 1024-token segments, 32 memory tokens). Only the configuration and the remote code were changed so that the model loads and runs correctly without access to the gated google/gemma-3-1b-it repository and without FlashAttention-2 (for example on a Colab T4, on CPU, or on Apple MPS).

from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("irodkin/ARMT-Gemma-3-1B-IT-8x1024-portable", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("irodkin/ARMT-Gemma-3-1B-IT-8x1024-portable")

Note: model.generate(...) returns only the newly generated tokens.

Changes relative to the original repository

config.json

• base_model_name replaced by an inline base_model_config (Gemma-3-1B text config), so the base architecture is built locally instead of downloading the gated google/gemma-3-1b-it checkpoint.
• armt_attn_implementation: "sdpa" (transformers consumes the reserved attn_implementation key).
• model_dtype, memory_dtype, dtype: float32 (cast with .to(...) or autocast as needed).

modeling_armt.py

• reverse_invert_attn_mask: compares float masks in floating point instead of casting finfo.min to an integer (undefined behaviour; it unmasked future tokens on Apple MPS), and treats boolean masks as "True = allowed" (the transformers SDPA path passes boolean masks for longer inputs; they were previously turned into an all-visible, non-causal mask).
• InnerLoopARMTForCausalLM: reads armt_attn_implementation, builds the base model with that kernel, and makes the layer wrappers use the kernel the base model actually uses (previously the wrappers could default to FlashAttention-2 masks while the base model ran SDPA).
• An inline dict base_model_config is built with AutoConfig.for_model(...), so version-specific conversions (e.g. Gemma-3 RoPE parameters in transformers 5) are applied.
• post_init() and get_input_embeddings / set_input_embeddings / get_output_embeddings added (required by transformers 5 and by PEFT).

Tested with transformers 4.57.1 and 5.1.0 on CPU and Apple MPS: identical outputs on both devices, no missing or unexpected weights, and LoRA (PEFT) backward passes.