Memory implementation for 1B

Hi can you please provide the config for 1B memory implementation I am trying to use content from 7B PK config but it doesnt work Thsi is the config I made dump_dir: runss/1bb_memory_model
name: "1b_pkplus_1024"
steps: 100_000
#probe_freq: 1000
seed: 777

optim:
lr: 3e-4
weight_decay: 0.01
warmup: 2000
lr_min_ratio: 0.0001
clip: 1.0
scheduler: cosine
annealing_step: 50000

distributed:
fsdp_type: full_shard
compile: true
model_dtype: bf16
matmul_allow_tf32: false
selective_activation_checkpointing: true
tp_size: 1
memory_parallel_size: 3

model:
dim: 2048
n_layers: 25
n_heads: 16
head_dim: 128
n_kv_heads: 4
rope_theta: 100_000
ffn_dim_multiplier: 1.3
multiple_of: 256
norm_eps: 1e-5
productkey_args:
is_enabled: True
layers: "6,13,20"
mem_share_values: True
mem_n_keys: 1024 # Retained, already a power of 2
value_fixed_lr: 0.0001
swilu_projection: true
mem_k_dim: 512 # Retained, already a power of 2

data:
root_dir: data/shuffled
sources:
user_dataset/arxiv: 1.0
batch_size: 24
prefetch_size: 128
seq_len: 1024
n_views: 2
load_async: true
add_bos: true
add_eos: true
tokenizer:
name: tiktoken
path: /home/gureshme/snap/firefox/memory/tokenizers/llama3/original/tokenizer.model

profiling:
run: true
mem_warmup: 100
mem_steps: 8
profile_warmup: 200
profile_steps: 8

checkpoint:
dump:
every: 5000
keep: 5

logging:
freq: 10

eval:
generator:
max_tokens: 8192
dtype: bf16
temperature: 1.0
top_p: 0.95
harness:
tasks:
- hellaswag
- piqa
- task: nq_open
num_fewshot: 5 and I keep getting this errror [rank0]: def embedding_bag_k(
[rank0]: out_ptr, # [B, dim]
[rank0]: indices_ptr, # [B, bag_size]
[rank0]: weight_ptr, # [n_keys**2, dim]
[rank0]: per_sample_weights, # [B, bag_size]
[rank0]: dim: tl.constexpr,
[rank0]: bag_size: tl.constexpr,
[rank0]: ):
[rank0]: out_idx = tl.program_id(axis=0).to(tl.int64)
[rank0]: out_value = tl.zeros([dim], dtype=tl.float32)
[rank0]: ^
[rank1]: Traceback (most recent call last):
[rank1]: File "/home/gureshme/.local/lib/python3.10/site-packages/triton/language/core.py", line 35, in wrapper
[rank1]: return fn(*args, **kwargs)
[rank1]: File "/home/gureshme/.local/lib/python3.10/site-packages/triton/language/core.py", line 1220, in full
[rank1]: shape = _shape_check_impl(shape)
[rank1]: File "/home/gureshme/.local/lib/python3.10/site-packages/triton/language/core.py", line 1205, in _shape_check_impl
[rank1]: raise ValueError(f"Shape element {i} must be a power of 2")
[rank1]: ValueError: Shape element 0 must be a power of 2