richard
/
learn_jax
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
learn_jax/parallel/gpt-neo-125m.json

854 lines
15 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
{
"transformer": {
"h": {
"0": {
"attn": {
"attention": {
"k_proj": {
"kernel": [
768,
768
]
},
"out_proj": {
"bias": [
768
],
"kernel": [
768,
768
]
},
"q_proj": {
"kernel": [
768,
768
]
},
"v_proj": {
"kernel": [
768,
768
]
}
}
},
"ln_1": {
"bias": [
768
],
"scale": [
768
]
},
"ln_2": {
"bias": [
768
],
"scale": [
768
]
},
"mlp": {
"c_fc": {
"bias": [
3072
],
"kernel": [
768,
3072
]
},
"c_proj": {
"bias": [
768
],
"kernel": [
3072,
768
]
}
}
},
"1": {
"attn": {
"attention": {
"k_proj": {
"kernel": [
768,
768
]
},
"out_proj": {
"bias": [
768
],
"kernel": [
768,
768
]
},
"q_proj": {
"kernel": [
768,
768
]
},
"v_proj": {
"kernel": [
768,
768
]
}
}
},
"ln_1": {
"bias": [
768
],
"scale": [
768
]
},
"ln_2": {
"bias": [
768
],
"scale": [
768
]
},
"mlp": {
"c_fc": {
"bias": [
3072
],
"kernel": [
768,
3072
]
},
"c_proj": {
"bias": [
768
],
"kernel": [
3072,
768
]
}
}
},
"10": {
"attn": {
"attention": {
"k_proj": {
"kernel": [
768,
768
]
},
"out_proj": {
"bias": [
768
],
"kernel": [
768,
768
]
},
"q_proj": {
"kernel": [
768,
768
]
},
"v_proj": {
"kernel": [
768,
768
]
}
}
},
"ln_1": {
"bias": [
768
],
"scale": [
768
]
},
"ln_2": {
"bias": [
768
],
"scale": [
768
]
},
"mlp": {
"c_fc": {
"bias": [
3072
],
"kernel": [
768,
3072
]
},
"c_proj": {
"bias": [
768
],
"kernel": [
3072,
768
]
}
}
},
"11": {
"attn": {
"attention": {
"k_proj": {
"kernel": [
768,
768
]
},
"out_proj": {
"bias": [
768
],
"kernel": [
768,
768
]
},
"q_proj": {
"kernel": [
768,
768
]
},
"v_proj": {
"kernel": [
768,
768
]
}
}
},
"ln_1": {
"bias": [
768
],
"scale": [
768
]
},
"ln_2": {
"bias": [
768
],
"scale": [
768
]
},
"mlp": {
"c_fc": {
"bias": [
3072
],
"kernel": [
768,
3072
]
},
"c_proj": {
"bias": [
768
],
"kernel": [
3072,
768
]
}
}
},
"2": {
"attn": {
"attention": {
"k_proj": {
"kernel": [
768,
768
]
},
"out_proj": {
"bias": [
768
],
"kernel": [
768,
768
]
},
"q_proj": {
"kernel": [
768,
768
]
},
"v_proj": {
"kernel": [
768,
768
]
}
}
},
"ln_1": {
"bias": [
768
],
"scale": [
768
]
},
"ln_2": {
"bias": [
768
],
"scale": [
768
]
},
"mlp": {
"c_fc": {
"bias": [
3072
],
"kernel": [
768,
3072
]
},
"c_proj": {
"bias": [
768
],
"kernel": [
3072,
768
]
}
}
},
"3": {
"attn": {
"attention": {
"k_proj": {
"kernel": [
768,
768
]
},
"out_proj": {
"bias": [
768
],
"kernel": [
768,
768
]
},
"q_proj": {
"kernel": [
768,
768
]
},
"v_proj": {
"kernel": [
768,
768
]
}
}
},
"ln_1": {
"bias": [
768
],
"scale": [
768
]
},
"ln_2": {
"bias": [
768
],
"scale": [
768
]
},
"mlp": {
"c_fc": {
"bias": [
3072
],
"kernel": [
768,
3072
]
},
"c_proj": {
"bias": [
768
],
"kernel": [
3072,
768
]
}
}
},
"4": {
"attn": {
"attention": {
"k_proj": {
"kernel": [
768,
768
]
},
"out_proj": {
"bias": [
768
],
"kernel": [
768,
768
]
},
"q_proj": {
"kernel": [
768,
768
]
},
"v_proj": {
"kernel": [
768,
768
]
}
}
},
"ln_1": {
"bias": [
768
],
"scale": [
768
]
},
"ln_2": {
"bias": [
768
],
"scale": [
768
]
},
"mlp": {
"c_fc": {
"bias": [
3072
],
"kernel": [
768,
3072
]
},
"c_proj": {
"bias": [
768
],
"kernel": [
3072,
768
]
}
}
},
"5": {
"attn": {
"attention": {
"k_proj": {
"kernel": [
768,
768
]
},
"out_proj": {
"bias": [
768
],
"kernel": [
768,
768
]
},
"q_proj": {
"kernel": [
768,
768
]
},
"v_proj": {
"kernel": [
768,
768
]
}
}
},
"ln_1": {
"bias": [
768
],
"scale": [
768
]
},
"ln_2": {
"bias": [
768
],
"scale": [
768
]
},
"mlp": {
"c_fc": {
"bias": [
3072
],
"kernel": [
768,
3072
]
},
"c_proj": {
"bias": [
768
],
"kernel": [
3072,
768
]
}
}
},
"6": {
"attn": {
"attention": {
"k_proj": {
"kernel": [
768,
768
]
},
"out_proj": {
"bias": [
768
],
"kernel": [
768,
768
]
},
"q_proj": {
"kernel": [
768,
768
]
},
"v_proj": {
"kernel": [
768,
768
]
}
}
},
"ln_1": {
"bias": [
768
],
"scale": [
768
]
},
"ln_2": {
"bias": [
768
],
"scale": [
768
]
},
"mlp": {
"c_fc": {
"bias": [
3072
],
"kernel": [
768,
3072
]
},
"c_proj": {
"bias": [
768
],
"kernel": [
3072,
768
]
}
}
},
"7": {
"attn": {
"attention": {
"k_proj": {
"kernel": [
768,
768
]
},
"out_proj": {
"bias": [
768
],
"kernel": [
768,
768
]
},
"q_proj": {
"kernel": [
768,
768
]
},
"v_proj": {
"kernel": [
768,
768
]
}
}
},
"ln_1": {
"bias": [
768
],
"scale": [
768
]
},
"ln_2": {
"bias": [
768
],
"scale": [
768
]
},
"mlp": {
"c_fc": {
"bias": [
3072
],
"kernel": [
768,
3072
]
},
"c_proj": {
"bias": [
768
],
"kernel": [
3072,
768
]
}
}
},
"8": {
"attn": {
"attention": {
"k_proj": {
"kernel": [
768,
768
]
},
"out_proj": {
"bias": [
768
],
"kernel": [
768,
768
]
},
"q_proj": {
"kernel": [
768,
768
]
},
"v_proj": {
"kernel": [
768,
768
]
}
}
},
"ln_1": {
"bias": [
768
],
"scale": [
768
]
},
"ln_2": {
"bias": [
768
],
"scale": [
768
]
},
"mlp": {
"c_fc": {
"bias": [
3072
],
"kernel": [
768,
3072
]
},
"c_proj": {
"bias": [
768
],
"kernel": [
3072,
768
]
}
}
},
"9": {
"attn": {
"attention": {
"k_proj": {
"kernel": [
768,
768
]
},
"out_proj": {
"bias": [
768
],
"kernel": [
768,
768
]
},
"q_proj": {
"kernel": [
768,
768
]
},
"v_proj": {
"kernel": [
768,
768
]
}
}
},
"ln_1": {
"bias": [
768
],
"scale": [
768
]
},
"ln_2": {
"bias": [
768
],
"scale": [
768
]
},
"mlp": {
"c_fc": {
"bias": [
3072
],
"kernel": [
768,
3072
]
},
"c_proj": {
"bias": [
768
],
"kernel": [
3072,
768
]
}
}
}
},
"ln_f": {
"bias": [
768
],
"scale": [
768
]
},
"wpe": {
"embedding": [
2048,
768
]
},
"wte": {
"embedding": [
50257,
768
]
}
}
}