Here is the code for layers of GPT-NeoX model and the code to load 20B checkpoint.
\nThe method _^_0_^_ in the layers load the checkpoints of that layer. The checkpoint loading helpers are on _^_1_^_
\n": "\u3053\u308c\u306f\u3001GPT-Neox\u30e2\u30c7\u30eb\u306e\u30ec\u30a4\u30e4\u30fc\u7528\u306e\u30b3\u30fc\u30c9\u306820B\u306e\u30c1\u30a7\u30c3\u30af\u30dd\u30a4\u30f3\u30c8\u3092\u30ed\u30fc\u30c9\u3059\u308b\u30b3\u30fc\u30c9\u3067\u3059\u3002
\n_^_0_^_\u30ec\u30a4\u30e4\u30fc\u5185\u306e\u30e1\u30bd\u30c3\u30c9\u306f\u3001\u305d\u306e\u30ec\u30a4\u30e4\u30fc\u306e\u30c1\u30a7\u30c3\u30af\u30dd\u30a4\u30f3\u30c8\u3092\u30ed\u30fc\u30c9\u3057\u307e\u3059\u3002\u30c1\u30a7\u30c3\u30af\u30dd\u30a4\u30f3\u30c8\u30ed\u30fc\u30c9\u30d8\u30eb\u30d1\u30fc\u304c\u30aa\u30f3\u306b\u306a\u3063\u3066\u3044\u307e\u3059 _^_1_^_
\n", "This is a standard embeddings layer with code to load the checkpoint.
\n": "\u3053\u308c\u306f\u3001\u30c1\u30a7\u30c3\u30af\u30dd\u30a4\u30f3\u30c8\u3092\u30ed\u30fc\u30c9\u3059\u308b\u30b3\u30fc\u30c9\u3092\u542b\u3080\u6a19\u6e96\u306e\u57cb\u3081\u8fbc\u307f\u30ec\u30a4\u30e4\u30fc\u3067\u3059\u3002
\n", "GPT-NeoX uses rotary positional embeddings (RoPE).
\nWE have annotated implementation of RoPE here with more notes the theory.
\n": "GPT-Neox\u306f\u56de\u8ee2\u5f0f\u30dd\u30b8\u30b7\u30e7\u30ca\u30eb\u30a8\u30f3\u30d9\u30c7\u30a3\u30f3\u30b0\uff08RoPE\uff09\u3092\u4f7f\u7528\u3057\u3066\u3044\u307e\u3059\u3002
\n\u3053\u3053\u3067\u306f\u3001RoPE \u306e\u5b9f\u88c5\u306b\u6ce8\u91c8\u3092\u4ed8\u3051\u3066\u3001\u7406\u8ad6\u306b\u95a2\u3059\u308b\u6ce8\u91c8\u3092\u4ed8\u3051\u307e\u3057\u305f\u3002
\n", "The layers are generated in the same order as checkpoints.
\nIt gives _^_0_^_ when a layer is not available; we use the layer indices as NeoX and there are two transformation layers we don't need in our implementation.
\n\u30ec\u30a4\u30e4\u30fc\u306f\u30c1\u30a7\u30c3\u30af\u30dd\u30a4\u30f3\u30c8\u3068\u540c\u3058\u9806\u5e8f\u3067\u751f\u6210\u3055\u308c\u307e\u3059\u3002
\n_^_0_^_\u30ec\u30a4\u30e4\u30fc\u304c\u4f7f\u7528\u3067\u304d\u306a\u3044\u5834\u5408\u306b\u8fd4\u3055\u308c\u307e\u3059\u3002\u30ec\u30a4\u30e4\u30fc\u30a4\u30f3\u30c7\u30c3\u30af\u30b9\u3092NeoX\u3068\u3057\u3066\u4f7f\u7528\u3057\u3001\u5b9f\u88c5\u306b\u306f\u5fc5\u8981\u306e\u306a\u3044\u5909\u63db\u30ec\u30a4\u30e4\u30fc\u304c2\u3064\u3042\u308a\u307e\u3059\u3002
\n_^_0_^_
\n": "_^_0_^_
\n", "Copying cached layers is faster than initializing new layers because it takes time to initialize parameters.
\nReturns the created layer or a copy of the cached layer
\u30ad\u30e3\u30c3\u30b7\u30e5\u3055\u308c\u305f\u30ec\u30a4\u30e4\u30fc\u306e\u30b3\u30d4\u30fc\u306f\u3001\u30d1\u30e9\u30e1\u30fc\u30bf\u30fc\u306e\u521d\u671f\u5316\u306b\u6642\u9593\u304c\u304b\u304b\u308b\u305f\u3081\u3001\u65b0\u3057\u3044\u30ec\u30a4\u30e4\u30fc\u3092\u521d\u671f\u5316\u3059\u308b\u3088\u308a\u3082\u9ad8\u901f\u3067\u3059\u3002
\n\u4f5c\u6210\u3055\u308c\u305f\u30ec\u30a4\u30e4\u30fc\u307e\u305f\u306f\u30ad\u30e3\u30c3\u30b7\u30e5\u3055\u308c\u305f\u30ec\u30a4\u30e4\u30fc\u306e\u30b3\u30d4\u30fc\u3092\u8fd4\u3057\u307e\u3059
We move the layer to the device and convert it to the correct data type
\nReturns the prepared layer
\u30ec\u30a4\u30e4\u30fc\u3092\u30c7\u30d0\u30a4\u30b9\u306b\u79fb\u52d5\u3057\u3001\u6b63\u3057\u3044\u30c7\u30fc\u30bf\u578b\u306b\u5909\u63db\u3057\u307e\u3059\u3002
\n\u6e96\u5099\u3057\u305f\u30ec\u30a4\u30e4\u30fc\u3092\u8fd4\u3057\u307e\u3059
\n": "\n", "\n
This function implements layer transformations after loading the checkpoint.
\nCurrently, it only applies the int8 quantization.
\nReturns the prepared layer
\u3053\u306e\u95a2\u6570\u306f\u3001\u30c1\u30a7\u30c3\u30af\u30dd\u30a4\u30f3\u30c8\u3092\u8aad\u307f\u8fbc\u3093\u3060\u5f8c\u306b\u30ec\u30a4\u30e4\u30fc\u5909\u63db\u3092\u5b9f\u88c5\u3057\u307e\u3059\u3002
\n\u73fe\u5728\u3001\u9069\u7528\u3055\u308c\u308b\u306e\u306f int8 \u91cf\u5b50\u5316\u306e\u307f\u3067\u3059\u3002
\n\u6e96\u5099\u3057\u305f\u30ec\u30a4\u30e4\u30fc\u3092\u8fd4\u3057\u307e\u3059
_^_0_^_
\n": "_^_0_^_
\n", "Code to load the checkpoint
\n": "\u30c1\u30a7\u30c3\u30af\u30dd\u30a4\u30f3\u30c8\u3092\u30ed\u30fc\u30c9\u3059\u308b\u30b3\u30fc\u30c9
\n", "Readout layer
\n": "\u8aad\u307f\u51fa\u3057\u5c64
\n", "\n": "\u30d5\u30e9\u30c3\u30b7\u30e5\u30fb\u30a2\u30c6\u30f3\u30b7\u30e7\u30f3
\n", "_^_0_^_
\n": "_^_0_^_
\n", "Add RoPE embeddings
\n": "RoPe \u57cb\u3081\u8fbc\u307f\u3092\u8ffd\u52a0
\n", "Add head dimension
\n": "\u982d\u90e8\u5bf8\u6cd5\u3092\u8ffd\u52a0
\n", "Add them and the residual connection
\n": "\u305d\u308c\u3089\u3068\u6b8b\u308a\u306e\u63a5\u7d9a\u3092\u8ffd\u52a0\u3057\u307e\u3059
\n", "Apply mask
\n": "\u30de\u30b9\u30af\u3092\u9069\u7528
\n", "Attention layer
\n": "\u30a2\u30c6\u30f3\u30b7\u30e7\u30f3\u30ec\u30a4\u30e4\u30fc
\n", "Attention output transform
\n": "\u30a2\u30c6\u30f3\u30b7\u30e7\u30f3\u51fa\u529b\u5909\u63db
\n", "Attention query, key and value transform
\n": "\u30a2\u30c6\u30f3\u30b7\u30e7\u30f3\u30af\u30a8\u30ea\u3001\u30ad\u30fc\u3001\u5024\u306e\u5909\u63db
\n", "Attention scaling factor
\n": "\u30a2\u30c6\u30f3\u30b7\u30e7\u30f3\u30b9\u30b1\u30fc\u30ea\u30f3\u30b0\u30d5\u30a1\u30af\u30bf\u30fc
\n", "Attention softmax
\n": "\u6ce8\u610f\u30bd\u30d5\u30c8\u30de\u30c3\u30af\u30b9
\n", "Attention softmax module
\n": "\u30a2\u30c6\u30f3\u30b7\u30e7\u30f3\u30bd\u30d5\u30c8\u30de\u30c3\u30af\u30b9\u30e2\u30b8\u30e5\u30fc\u30eb
\n", "Base for _^_0_^_
\n": "\u306e\u30d9\u30fc\u30b9 _^_0_^_
\n", "Cache _^_0_^_ and _^_1_^_
\n": "_^_0_^_\u30ad\u30e3\u30c3\u30b7\u30e5\u3068 _^_1_^_
\n", "Cache them
\n": "\u305d\u308c\u3089\u3092\u30ad\u30e3\u30c3\u30b7\u30e5\u3059\u308b
\n", "Calculate _^_0_^_ and _^_1_^_ in fp32
\n": "_^_0_^_\u8a08\u7b97\u3057\u3066 _^_1_^_ fp32 \u3067
\n", "Concatenate so that for row _^_0_^_ we have
\n_^_1_^_
\n": "\u884c\u304c\u6b21\u306e\u3088\u3046\u306b\u306a\u308b\u3088\u3046\u306b\u9023\u7d50\u3057\u307e\u3059 _^_0_^_
\n_^_1_^_
\n", "Concatenate the past
\n": "\u904e\u53bb\u3092\u9023\u7d50\u3059\u308b
\n", "Concatenate with features that didn't get RoPE embeddings
\n": "RoPe \u57cb\u3081\u8fbc\u307f\u306b\u5bfe\u5fdc\u3057\u3066\u3044\u306a\u304b\u3063\u305f\u6a5f\u80fd\u3068\u306e\u9023\u643a
\n", "Contraction linear layer
\n": "\u53ce\u7e2e\u7dda\u72b6\u5c64
\n", "Convert the linear layers
\n": "\u7dda\u5f62\u30ec\u30a4\u30e4\u30fc\u306e\u5909\u63db
\n", "Convert to fp32 if the current dtype is fp16
\n": "\u73fe\u5728\u306e dtype \u304c fp16 \u306e\u5834\u5408\u306f fp32 \u306b\u5909\u63db
\n", "Create mask
\n": "\u30de\u30b9\u30af\u4f5c\u6210
\n", "Disable auto-casting to fp16 for attention computation
\n": "\u30a2\u30c6\u30f3\u30b7\u30e7\u30f3\u8a08\u7b97\u306e fp16 \u3078\u306e\u81ea\u52d5\u30ad\u30e3\u30b9\u30c8\u3092\u7121\u52b9\u306b\u3059\u308b
\n", "Do not cast for bfloat
\n": "bfloat\u306b\u306f\u30ad\u30e3\u30b9\u30c8\u3057\u306a\u3044\u3067\u304f\u3060\u3055\u3044
\n", "Embedding layer
\n": "\u57cb\u3081\u8fbc\u307f\u30ec\u30a4\u30e4\u30fc
\n", "Expansion linear layer
\n": "\u62e1\u5f35\u30ea\u30cb\u30a2\u30ec\u30a4\u30e4\u30fc
\n", "FFN first transform
\n": "FFN \u30d5\u30a1\u30fc\u30b9\u30c8\u30c8\u30e9\u30f3\u30b9\u30d5\u30a9\u30fc\u30e0
\n", "FFN layer
\n": "FFN \u30ec\u30a4\u30e4\u30fc
\n", "FFN second transform
\n": "FFN 2 \u756a\u76ee\u306e\u30c8\u30e9\u30f3\u30b9\u30d5\u30a9\u30fc\u30e0
\n", "Final linear layer
\n": "\u6700\u7d42\u7dda\u5f62\u30ec\u30a4\u30e4\u30fc
\n", "Final normalization layer
\n": "\u6700\u7d42\u6b63\u898f\u5316\u30ec\u30a4\u30e4\u30fc
\n", "GELU activation
\n": "GELU \u30a2\u30af\u30c6\u30a3\u30d9\u30fc\u30b7\u30e7\u30f3
\n", "Get attention weighted values
\n": "\u30a2\u30c6\u30f3\u30b7\u30e7\u30f3\u52a0\u91cd\u5024\u3092\u53d6\u5f97
\n", "Get causal mask
\n": "\u30ab\u30b8\u30e5\u30a2\u30eb\u30de\u30b9\u30af\u3092\u30b2\u30c3\u30c8
\n", "Get default values if not specified
\n": "\u6307\u5b9a\u3057\u306a\u3044\u5834\u5408\u306f\u30c7\u30d5\u30a9\u30eb\u30c8\u5024\u3092\u53d6\u5f97
\n", "Get position indexes _^_0_^_
\n": "\u4f4d\u7f6e\u30a4\u30f3\u30c7\u30c3\u30af\u30b9\u3092\u53d6\u5f97 _^_0_^_
\n", "Get query, key and value embeddings (all concatenated). The last dimension size will change from n_hidden -> _^_0_^_
\n": "\u30af\u30a8\u30ea\u3001\u30ad\u30fc\u3001\u5024\u306e\u57cb\u3081\u8fbc\u307f (\u3059\u3079\u3066\u9023\u7d50) \u3092\u53d6\u5f97\u3057\u307e\u3059\u3002\u6700\u5f8c\u306e\u30c7\u30a3\u30e1\u30f3\u30b7\u30e7\u30f3\u30b5\u30a4\u30ba\u304c n_hidden \u304b\u3089\u5909\u66f4\u3055\u308c\u307e\u3059
-> _^_0_^_\n", "Get the actual sequence length
\n": "\u5b9f\u969b\u306e\u30b7\u30fc\u30b1\u30f3\u30b9\u9577\u3092\u53d6\u5f97
\n", "Get the past keys and values. These will have shape _^_0_^_
\n": "\u904e\u53bb\u306e\u30ad\u30fc\u3068\u5024\u3092\u53d6\u5f97\u3057\u307e\u3059\u3002\u3053\u308c\u3089\u306f\u5f62\u306b\u306a\u308a\u307e\u3059 _^_0_^_
\n", "Get the sin and cos values from the cache
\n": "\u30ad\u30e3\u30c3\u30b7\u30e5\u304b\u3089 sin \u3068 cos \u306e\u5024\u3092\u53d6\u5f97
\n", "Get the state id's. We use to retrieve previous states and store the next states
\n": "\u30b9\u30c6\u30fc\u30c8 ID \u3092\u53d6\u5f97\u3057\u307e\u3059\u3002\u524d\u306e\u30b9\u30c6\u30fc\u30c8\u3092\u53d6\u5f97\u3057\u305f\u308a\u3001\u6b21\u306e\u30b9\u30c6\u30fc\u30c8\u3092\u4fdd\u5b58\u3057\u305f\u308a\u3059\u308b\u306e\u306b\u4f7f\u3044\u307e\u3059\u3002
\n", "If there's cache
\n": "\u30ad\u30e3\u30c3\u30b7\u30e5\u304c\u3042\u308b\u5834\u5408
\n", "If we are caching the states of previous tokens
\n": "\u4ee5\u524d\u306e\u30c8\u30fc\u30af\u30f3\u306e\u72b6\u614b\u3092\u30ad\u30e3\u30c3\u30b7\u30e5\u3059\u308b\u5834\u5408
\n", "Initialize _^_0_^_
\n": "[\u521d\u671f\u5316] _^_0_^_
\n", "Initialize _^_0_^_ and _^_1_^_ cache
\n": "_^_0_^_\u521d\u671f\u5316\u3068\u30ad\u30e3\u30c3\u30b7\u30e5 _^_1_^_
\n", "Layer norm before FFN
\n": "FFN \u524d\u306e\u30ec\u30a4\u30e4\u30fc\u30ce\u30eb\u30e0
\n", "Layer norm before attention
\n": "\u6ce8\u76ee\u3055\u308c\u308b\u524d\u306e\u30ec\u30a4\u30e4\u30fc\u30ce\u30eb\u30e0
\n", "Layer normalization before FFN
\n": "FFN \u524d\u306e\u30ec\u30a4\u30e4\u30fc\u6b63\u898f\u5316
\n", "Layer normalization before attention
\n": "\u6ce8\u610f\u524d\u306e\u30ec\u30a4\u30e4\u30fc\u6b63\u898f\u5316
\n", "Linear layer for query, key and value
\n": "\u30af\u30a8\u30ea\u3001\u30ad\u30fc\u3001\u5024\u306e\u7dda\u5f62\u30ec\u30a4\u30e4\u30fc
\n", "NeoX runs attention and feedforward network in parallel
\n": "NeoX\u306f\u30a2\u30c6\u30f3\u30b7\u30e7\u30f3\u30cd\u30c3\u30c8\u30ef\u30fc\u30af\u3068\u30d5\u30a3\u30fc\u30c9\u30d5\u30a9\u30ef\u30fc\u30c9\u30cd\u30c3\u30c8\u30ef\u30fc\u30af\u3092\u4e26\u884c\u3057\u3066\u5b9f\u884c\u3057\u307e\u3059
\n", "No cache - simply add RoPE embeddings
\n": "\u30ad\u30e3\u30c3\u30b7\u30e5\u306a\u3057-RoPE \u57cb\u3081\u8fbc\u307f\u3092\u8ffd\u52a0\u3059\u308b\u3060\u3051
\n", "Number of features for RoPE
\n": "RoPE \u306e\u6a5f\u80fd\u306e\u6570
\n", "Number of features per head
\n": "\u30d8\u30c3\u30c9\u3042\u305f\u308a\u306e\u6a5f\u80fd\u6570
\n", "Offset of the current embeddings
\n": "\u73fe\u5728\u306e\u57cb\u3081\u8fbc\u307f\u306e\u30aa\u30d5\u30bb\u30c3\u30c8
\n", "Only convert the linear layers in the transformer layers
\n": "\u30c8\u30e9\u30f3\u30b9\u30ec\u30a4\u30e4\u30fc\u306e\u7dda\u5f62\u30ec\u30a4\u30e4\u30fc\u306e\u307f\u3092\u5909\u63db\u3057\u307e\u3059
\n", "Otherwise, use normal attention
\n": "\u305d\u308c\u4ee5\u5916\u306e\u5834\u5408\u306f\u3001\u901a\u5e38\u306e\u6ce8\u610f\u3092\u6255\u3063\u3066\u304f\u3060\u3055\u3044
\n", "Query and key lengths
\n": "\u30af\u30a8\u30ea\u3068\u30ad\u30fc\u306e\u9577\u3055
\n", "Readout layer
\n": "\u8aad\u307f\u51fa\u3057\u5c64
\n", "Reshape from _^_0_^_batch_size, seq_len, n_hidden`
\n": "_^_0_^_\u30d0\u30c3\u30c1\u30b5\u30a4\u30ba\u3001\u30b7\u30fc\u30b1\u30f3\u30b9\u756a\u53f7\u3001n_hidden `\u304b\u3089\u5f62\u72b6\u3092\u5909\u66f4
\n", "Residual connection
\n": "\u6b8b\u7559\u63a5\u7d9a
\n", "Return from cache
\n": "\u30ad\u30e3\u30c3\u30b7\u30e5\u304b\u3089\u623b\u308b
\n", "RoPE embedding module
\n": "RoPE \u57cb\u3081\u8fbc\u307f\u30e2\u30b8\u30e5\u30fc\u30eb
\n", "RoPE embeddings
\n_^_0_^_for _^_1_^_
\n": "\u30ed\u30fc\u30d7\u57cb\u3081\u8fbc\u307f
\n_^_0_^_\u306b\u3068\u3063\u3066 _^_1_^_
\n", "Save the current state
\n": "\u73fe\u5728\u306e\u72b6\u614b\u3092\u4fdd\u5b58\u3059\u308b
\n", "Scale attention
\n": "\u30b9\u30b1\u30fc\u30eb\u30a2\u30c6\u30f3\u30b7\u30e7\u30f3
\n", "Skip if not using int8 quantization
\n": "int8 \u91cf\u5b50\u5316\u3092\u4f7f\u7528\u3057\u306a\u3044\u5834\u5408\u306f\u30b9\u30ad\u30c3\u30d7
\n", "Split into heads by changing the shape to _^_0_^_
\n": "\u5f62\u72b6\u3092\u4ee5\u4e0b\u306e\u3088\u3046\u306b\u5909\u66f4\u3057\u3066\u982d\u90e8\u306b\u5206\u5272\u3057\u307e\u3059 _^_0_^_
\n", "Split into query, key and value each of shape _^_0_^_
\n": "\u5f62\u72b6\u3054\u3068\u306b\u30af\u30a8\u30ea\u3001\u30ad\u30fc\u3001\u5024\u306b\u5206\u5272 _^_0_^_
\n", "Split the features. We apply RoPE to only _^_0_^_ features
\n": "\u6a5f\u80fd\u3092\u5206\u5272\u3057\u3066\u304f\u3060\u3055\u3044\u3002RoPE _^_0_^_ \u306f\u6a5f\u80fd\u306b\u306e\u307f\u9069\u7528\u3055\u308c\u307e\u3059
\n", "Stack them into shape _^_0_^_
\n": "\u305d\u308c\u3089\u3092\u7a4d\u307f\u91cd\u306d\u3066\u5f62\u3092\u6574\u3048\u308b _^_0_^_
\n", "The output is of shape _^_0_^_
\n": "\u51fa\u529b\u306f\u6574\u5f62\u3057\u3066\u3044\u307e\u3059 _^_0_^_
\n", "To cache causal mask
\n": "\u56e0\u679c\u30de\u30b9\u30af\u3092\u30ad\u30e3\u30c3\u30b7\u30e5\u3059\u308b\u306b\u306f
\n", "To store _^_0_^_ for the features
\n": "_^_0_^_\u6a5f\u80fd\u7528\u306b\u4fdd\u5b58\u3059\u308b\u306b\u306f
\n", "Transformer layer
\n": "\u5909\u5727\u5668\u5c64
\n", "Transformer layers
\n": "\u30c8\u30e9\u30f3\u30b9\u30d5\u30a9\u30fc\u30de\u30fc\u5c64
\n", "Use _^_0_^_ defined in utilities.
\n": "_^_0_^_\u30e6\u30fc\u30c6\u30a3\u30ea\u30c6\u30a3\u3067\u5b9a\u7fa9\u3055\u308c\u3066\u3044\u308b\u7528\u9014\u3002
\n", "Use flash attention
\n": "\u30d5\u30e9\u30c3\u30b7\u30e5\u30a2\u30c6\u30f3\u30b7\u30e7\u30f3\u3092\u4f7f\u3046
\n", "Out implementation doesn't include dropout.
\n": "\u30a2\u30a6\u30c8\u306e\u5b9f\u88c5\u306b\u306f\u30c9\u30ed\u30c3\u30d7\u30a2\u30a6\u30c8\u306f\u542b\u307e\u308c\u3066\u3044\u307e\u305b\u3093\u3002
\n", "