A test for using symbolic shapes for minformer.

[gnecula]

The actual minformer code is from https://github.com/sholtodouglas/minformer/blob/main/minformer/model_test.py with as few modifications as possible.

The modifications needed to get the symbolic evaluations to work are marked with TODO.

To run the examples:

pytest --verbosity 1 -s tests/minformer_symbolic_shape_test.py

The test output:

tests/minformer_symbolic_shape_test.py::MinformerSymbolicShapesTest::test_inference_causal=False
weights shapes: Weights(layers=[Layer(q=(D, N, H), k=(D, N, H), v=(D, N, H), proj=(N, H, D), w1=(D, D*F), w2=(D*F, D), attn_gamma=(D,), ffn_gamma=(D,), q_gamma=(H,), k_gamma=(H,)), Layer(q=(D, N, H), k=(D, N, H), v=(D, N, H), proj=(N, H, D), w1=(D, D*F), w2=(D*F, D), attn_gamma=(D,), ffn_gamma=(D,), q_gamma=(H,), k_gamma=(H,)), Layer(q=(D, N, H), k=(D, N, H), v=(D, N, H), proj=(N, H, D), w1=(D, D*F), w2=(D*F, D), attn_gamma=(D,), ffn_gamma=(D,), q_gamma=(H,), k_gamma=(H,)), Layer(q=(D, N, H), k=(D, N, H), v=(D, N, H), proj=(N, H, D), w1=(D, D*F), w2=(D*F, D), attn_gamma=(D,), ffn_gamma=(D,), q_gamma=(H,), k_gamma=(H,))], embedding=(V, D), gamma_final=(D,))
prefill_cache shapes: KVCache(k=[(1, N, T, H), (1, N, T, H), (1, N, T, H), (1, N, T, H)], v=[(1, N, T, H), (1, N, T, H), (1, N, T, H), (1, N, T, H)], lengths=(1,))
chunk_a shapes: (1, S)
segment_ids shapes: (1, S)
logits shapes: (1, S, V)
PASSED
tests/minformer_symbolic_shape_test.py::MinformerSymbolicShapesTest::test_inference_causal=True
weights shapes: Weights(layers=[Layer(q=(D, N, H), k=(D, N, H), v=(D, N, H), proj=(N, H, D), w1=(D, D*F), w2=(D*F, D), attn_gamma=(D,), ffn_gamma=(D,), q_gamma=(H,), k_gamma=(H,)), Layer(q=(D, N, H), k=(D, N, H), v=(D, N, H), proj=(N, H, D), w1=(D, D*F), w2=(D*F, D), attn_gamma=(D,), ffn_gamma=(D,), q_gamma=(H,), k_gamma=(H,)), Layer(q=(D, N, H), k=(D, N, H), v=(D, N, H), proj=(N, H, D), w1=(D, D*F), w2=(D*F, D), attn_gamma=(D,), ffn_gamma=(D,), q_gamma=(H,), k_gamma=(H,)), Layer(q=(D, N, H), k=(D, N, H), v=(D, N, H), proj=(N, H, D), w1=(D, D*F), w2=(D*F, D), attn_gamma=(D,), ffn_gamma=(D,), q_gamma=(H,), k_gamma=(H,))], embedding=(V, D), gamma_final=(D,))
prefill_cache shapes: KVCache(k=[(1, N, T, H), (1, N, T, H), (1, N, T, H), (1, N, T, H)], v=[(1, N, T, H), (1, N, T, H), (1, N, T, H), (1, N, T, H)], lengths=(1,))
chunk_a shapes: (1, S)
segment_ids shapes: (1, S)
logits shapes: (1, S, V)
PASSED
tests/minformer_symbolic_shape_test.py::MinformerSymbolicShapesTest::test_overtrain_and_sample_simple_sequence SKIPPED (Works only on TPU)
tests/minformer_symbolic_shape_test.py::MinformerSymbolicShapesTest::test_training_use_attn_kernel=False
weights shapes: Weights(layers=[Layer(q=(D, N, H), k=(D, N, H), v=(D, N, H), proj=(N, H, D), w1=(D, D*F), w2=(D*F, D), attn_gamma=(D,), ffn_gamma=(D,), q_gamma=(H,), k_gamma=(H,)), Layer(q=(D, N, H), k=(D, N, H), v=(D, N, H), proj=(N, H, D), w1=(D, D*F), w2=(D*F, D), attn_gamma=(D,), ffn_gamma=(D,), q_gamma=(H,), k_gamma=(H,)), Layer(q=(D, N, H), k=(D, N, H), v=(D, N, H), proj=(N, H, D), w1=(D, D*F), w2=(D*F, D), attn_gamma=(D,), ffn_gamma=(D,), q_gamma=(H,), k_gamma=(H,)), Layer(q=(D, N, H), k=(D, N, H), v=(D, N, H), proj=(N, H, D), w1=(D, D*F), w2=(D*F, D), attn_gamma=(D,), ffn_gamma=(D,), q_gamma=(H,), k_gamma=(H,))], embedding=(V, D), gamma_final=(D,))
opt_state shapes: Weights(layers=[Layer(q=((D, N, H), (D, N, H), (D, N, H)), k=((D, N, H), (D, N, H), (D, N, H)), v=((D, N, H), (D, N, H), (D, N, H)), proj=((N, H, D), (N, H, D), (N, H, D)), w1=((D, D*F), (D, D*F), (D, D*F)), w2=((D*F, D), (D*F, D), (D*F, D)), attn_gamma=((D,), (D,), (D,)), ffn_gamma=((D,), (D,), (D,)), q_gamma=((H,), (H,), (H,)), k_gamma=((H,), (H,), (H,))), Layer(q=((D, N, H), (D, N, H), (D, N, H)), k=((D, N, H), (D, N, H), (D, N, H)), v=((D, N, H), (D, N, H), (D, N, H)), proj=((N, H, D), (N, H, D), (N, H, D)), w1=((D, D*F), (D, D*F), (D, D*F)), w2=((D*F, D), (D*F, D), (D*F, D)), attn_gamma=((D,), (D,), (D,)), ffn_gamma=((D,), (D,), (D,)), q_gamma=((H,), (H,), (H,)), k_gamma=((H,), (H,), (H,))), Layer(q=((D, N, H), (D, N, H), (D, N, H)), k=((D, N, H), (D, N, H), (D, N, H)), v=((D, N, H), (D, N, H), (D, N, H)), proj=((N, H, D), (N, H, D), (N, H, D)), w1=((D, D*F), (D, D*F), (D, D*F)), w2=((D*F, D), (D*F, D), (D*F, D)), attn_gamma=((D,), (D,), (D,)), ffn_gamma=((D,), (D,), (D,)), q_gamma=((H,), (H,), (H,)), k_gamma=((H,), (H,), (H,))), Layer(q=((D, N, H), (D, N, H), (D, N, H)), k=((D, N, H), (D, N, H), (D, N, H)), v=((D, N, H), (D, N, H), (D, N, H)), proj=((N, H, D), (N, H, D), (N, H, D)), w1=((D, D*F), (D, D*F), (D, D*F)), w2=((D*F, D), (D*F, D), (D*F, D)), attn_gamma=((D,), (D,), (D,)), ffn_gamma=((D,), (D,), (D,)), q_gamma=((H,), (H,), (H,)), k_gamma=((H,), (H,), (H,)))], embedding=((V, D), (V, D), (V, D)), gamma_final=((D,), (D,), (D,)))
prefill_cache shapes: KVCache(k=[(B, N, T, H), (B, N, T, H), (B, N, T, H), (B, N, T, H)], v=[(B, N, T, H), (B, N, T, H), (B, N, T, H), (B, N, T, H)], lengths=(B,))
loss shapes: ()
cache shapes: KVCache(k=[(1, N, T, H), (1, N, T, H), (1, N, T, H), (1, N, T, H)], v=[(1, N, T, H), (1, N, T, H), (1, N, T, H), (1, N, T, H)], lengths=(1,))
PASSED
tests/minformer_symbolic_shape_test.py::MinformerSymbolicShapesTest::test_training_use_attn_kernel=True
weights shapes: Weights(layers=[Layer(q=(D, N, 8), k=(D, N, 8), v=(D, N, 8), proj=(N, 8, D), w1=(D, D*F), w2=(D*F, D), attn_gamma=(D,), ffn_gamma=(D,), q_gamma=(8,), k_gamma=(8,)), Layer(q=(D, N, 8), k=(D, N, 8), v=(D, N, 8), proj=(N, 8, D), w1=(D, D*F), w2=(D*F, D), attn_gamma=(D,), ffn_gamma=(D,), q_gamma=(8,), k_gamma=(8,)), Layer(q=(D, N, 8), k=(D, N, 8), v=(D, N, 8), proj=(N, 8, D), w1=(D, D*F), w2=(D*F, D), attn_gamma=(D,), ffn_gamma=(D,), q_gamma=(8,), k_gamma=(8,)), Layer(q=(D, N, 8), k=(D, N, 8), v=(D, N, 8), proj=(N, 8, D), w1=(D, D*F), w2=(D*F, D), attn_gamma=(D,), ffn_gamma=(D,), q_gamma=(8,), k_gamma=(8,))], embedding=(V, D), gamma_final=(D,))
opt_state shapes: Weights(layers=[Layer(q=((D, N, 8), (D, N, 8), (D, N, 8)), k=((D, N, 8), (D, N, 8), (D, N, 8)), v=((D, N, 8), (D, N, 8), (D, N, 8)), proj=((N, 8, D), (N, 8, D), (N, 8, D)), w1=((D, D*F), (D, D*F), (D, D*F)), w2=((D*F, D), (D*F, D), (D*F, D)), attn_gamma=((D,), (D,), (D,)), ffn_gamma=((D,), (D,), (D,)), q_gamma=((8,), (8,), (8,)), k_gamma=((8,), (8,), (8,))), Layer(q=((D, N, 8), (D, N, 8), (D, N, 8)), k=((D, N, 8), (D, N, 8), (D, N, 8)), v=((D, N, 8), (D, N, 8), (D, N, 8)), proj=((N, 8, D), (N, 8, D), (N, 8, D)), w1=((D, D*F), (D, D*F), (D, D*F)), w2=((D*F, D), (D*F, D), (D*F, D)), attn_gamma=((D,), (D,), (D,)), ffn_gamma=((D,), (D,), (D,)), q_gamma=((8,), (8,), (8,)), k_gamma=((8,), (8,), (8,))), Layer(q=((D, N, 8), (D, N, 8), (D, N, 8)), k=((D, N, 8), (D, N, 8), (D, N, 8)), v=((D, N, 8), (D, N, 8), (D, N, 8)), proj=((N, 8, D), (N, 8, D), (N, 8, D)), w1=((D, D*F), (D, D*F), (D, D*F)), w2=((D*F, D), (D*F, D), (D*F, D)), attn_gamma=((D,), (D,), (D,)), ffn_gamma=((D,), (D,), (D,)), q_gamma=((8,), (8,), (8,)), k_gamma=((8,), (8,), (8,))), Layer(q=((D, N, 8), (D, N, 8), (D, N, 8)), k=((D, N, 8), (D, N, 8), (D, N, 8)), v=((D, N, 8), (D, N, 8), (D, N, 8)), proj=((N, 8, D), (N, 8, D), (N, 8, D)), w1=((D, D*F), (D, D*F), (D, D*F)), w2=((D*F, D), (D*F, D), (D*F, D)), attn_gamma=((D,), (D,), (D,)), ffn_gamma=((D,), (D,), (D,)), q_gamma=((8,), (8,), (8,)), k_gamma=((8,), (8,), (8,)))], embedding=((V, D), (V, D), (V, D)), gamma_final=((D,), (D,), (D,)))
prefill_cache shapes: KVCache(k=[(B, N, T, 8), (B, N, T, 8), (B, N, T, 8), (B, N, T, 8)], v=[(B, N, T, 8), (B, N, T, 8), (B, N, T, 8), (B, N, T, 8)], lengths=(B,))
loss shapes: ()
cache shapes: KVCache(k=[(1, N, T, 8), (1, N, T, 8), (1, N, T, 8), (1, N, T, 8)], v=[(1, N, T, 8), (1, N, T, 8), (1, N, T, 8), (1, N, T, 8)], lengths=(1,))
PASSED