richard
/
learn_jax
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
learn_jax/parallel/t5.json

1826 lines
37 KiB
JSON
Raw Normal View History

Feat: t5_jax_simple_parallel implements a working example of fsdp
2024-09-20 23:42:51 +09:00
{
"decoder": {
"block": {
"0": {
"layer": {
"0": {
"SelfAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"relative_attention_bias": {
"embedding": [
32,
12
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"1": {
"EncDecAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"2": {
"DenseReluDense": {
"wi": {
"kernel": [
768,
3072
]
},
"wo": {
"kernel": [
3072,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
}
}
},
"1": {
"layer": {
"0": {
"SelfAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"1": {
"EncDecAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"2": {
"DenseReluDense": {
"wi": {
"kernel": [
768,
3072
]
},
"wo": {
"kernel": [
3072,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
}
}
},
"10": {
"layer": {
"0": {
"SelfAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"1": {
"EncDecAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"2": {
"DenseReluDense": {
"wi": {
"kernel": [
768,
3072
]
},
"wo": {
"kernel": [
3072,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
}
}
},
"11": {
"layer": {
"0": {
"SelfAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"1": {
"EncDecAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"2": {
"DenseReluDense": {
"wi": {
"kernel": [
768,
3072
]
},
"wo": {
"kernel": [
3072,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
}
}
},
"2": {
"layer": {
"0": {
"SelfAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"1": {
"EncDecAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"2": {
"DenseReluDense": {
"wi": {
"kernel": [
768,
3072
]
},
"wo": {
"kernel": [
3072,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
}
}
},
"3": {
"layer": {
"0": {
"SelfAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"1": {
"EncDecAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"2": {
"DenseReluDense": {
"wi": {
"kernel": [
768,
3072
]
},
"wo": {
"kernel": [
3072,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
}
}
},
"4": {
"layer": {
"0": {
"SelfAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"1": {
"EncDecAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"2": {
"DenseReluDense": {
"wi": {
"kernel": [
768,
3072
]
},
"wo": {
"kernel": [
3072,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
}
}
},
"5": {
"layer": {
"0": {
"SelfAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"1": {
"EncDecAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"2": {
"DenseReluDense": {
"wi": {
"kernel": [
768,
3072
]
},
"wo": {
"kernel": [
3072,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
}
}
},
"6": {
"layer": {
"0": {
"SelfAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"1": {
"EncDecAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"2": {
"DenseReluDense": {
"wi": {
"kernel": [
768,
3072
]
},
"wo": {
"kernel": [
3072,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
}
}
},
"7": {
"layer": {
"0": {
"SelfAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"1": {
"EncDecAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"2": {
"DenseReluDense": {
"wi": {
"kernel": [
768,
3072
]
},
"wo": {
"kernel": [
3072,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
}
}
},
"8": {
"layer": {
"0": {
"SelfAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"1": {
"EncDecAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"2": {
"DenseReluDense": {
"wi": {
"kernel": [
768,
3072
]
},
"wo": {
"kernel": [
3072,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
}
}
},
"9": {
"layer": {
"0": {
"SelfAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"1": {
"EncDecAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"2": {
"DenseReluDense": {
"wi": {
"kernel": [
768,
3072
]
},
"wo": {
"kernel": [
3072,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
}
}
}
},
"final_layer_norm": {
"weight": [
768
]
}
},
"encoder": {
"block": {
"0": {
"layer": {
"0": {
"SelfAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"relative_attention_bias": {
"embedding": [
32,
12
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"1": {
"DenseReluDense": {
"wi": {
"kernel": [
768,
3072
]
},
"wo": {
"kernel": [
3072,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
}
}
},
"1": {
"layer": {
"0": {
"SelfAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"1": {
"DenseReluDense": {
"wi": {
"kernel": [
768,
3072
]
},
"wo": {
"kernel": [
3072,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
}
}
},
"10": {
"layer": {
"0": {
"SelfAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"1": {
"DenseReluDense": {
"wi": {
"kernel": [
768,
3072
]
},
"wo": {
"kernel": [
3072,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
}
}
},
"11": {
"layer": {
"0": {
"SelfAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"1": {
"DenseReluDense": {
"wi": {
"kernel": [
768,
3072
]
},
"wo": {
"kernel": [
3072,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
}
}
},
"2": {
"layer": {
"0": {
"SelfAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"1": {
"DenseReluDense": {
"wi": {
"kernel": [
768,
3072
]
},
"wo": {
"kernel": [
3072,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
}
}
},
"3": {
"layer": {
"0": {
"SelfAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"1": {
"DenseReluDense": {
"wi": {
"kernel": [
768,
3072
]
},
"wo": {
"kernel": [
3072,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
}
}
},
"4": {
"layer": {
"0": {
"SelfAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"1": {
"DenseReluDense": {
"wi": {
"kernel": [
768,
3072
]
},
"wo": {
"kernel": [
3072,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
}
}
},
"5": {
"layer": {
"0": {
"SelfAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"1": {
"DenseReluDense": {
"wi": {
"kernel": [
768,
3072
]
},
"wo": {
"kernel": [
3072,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
}
}
},
"6": {
"layer": {
"0": {
"SelfAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"1": {
"DenseReluDense": {
"wi": {
"kernel": [
768,
3072
]
},
"wo": {
"kernel": [
3072,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
}
}
},
"7": {
"layer": {
"0": {
"SelfAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"1": {
"DenseReluDense": {
"wi": {
"kernel": [
768,
3072
]
},
"wo": {
"kernel": [
3072,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
}
}
},
"8": {
"layer": {
"0": {
"SelfAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"1": {
"DenseReluDense": {
"wi": {
"kernel": [
768,
3072
]
},
"wo": {
"kernel": [
3072,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
}
}
},
"9": {
"layer": {
"0": {
"SelfAttention": {
"k": {
"kernel": [
768,
768
]
},
"o": {
"kernel": [
768,
768
]
},
"q": {
"kernel": [
768,
768
]
},
"v": {
"kernel": [
768,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
},
"1": {
"DenseReluDense": {
"wi": {
"kernel": [
768,
3072
]
},
"wo": {
"kernel": [
3072,
768
]
}
},
"layer_norm": {
"weight": [
768
]
}
}
}
}
},
"final_layer_norm": {
"weight": [
768
]
}
},
"shared": {
"embedding": [
32128,
768
]
}
}