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": "\u4ee5\u4e0b\u662f GPT-NEOX \u6a21\u578b\u5c42\u7684\u4ee3\u7801\u548c\u52a0\u8f7d 20B \u68c0\u67e5\u70b9\u7684\u4ee3\u7801\u3002
\n\u56fe\u5c42_^_0_^_\u4e2d\u7684\u65b9\u6cd5\u52a0\u8f7d\u8be5\u5c42\u7684\u68c0\u67e5\u70b9\u3002\u68c0\u67e5\u70b9\u52a0\u8f7d\u52a9\u624b\u5df2\u542f\u7528 _^_1_^_
\n", "This is a standard embeddings layer with code to load the checkpoint.
\n": "\u8fd9\u662f\u4e00\u4e2a\u6807\u51c6\u7684\u5d4c\u5165\u5c42\uff0c\u5176\u4e2d\u5305\u542b\u7528\u4e8e\u52a0\u8f7d\u68c0\u67e5\u70b9\u7684\u4ee3\u7801\u3002
\n", "GPT-NeoX uses rotary positional embeddings (RoPE).
\nWE have annotated implementation of RoPE here with more notes the theory.
\n": "GPT-NEOX \u4f7f\u7528\u65cb\u8f6c\u4f4d\u7f6e\u5d4c\u5165\uff08RoP\uff09\u3002
\n\u6211\u4eec\u5728\u8fd9\u91cc\u6ce8\u91ca\u4e86 RoPe \u7684\u5b9e\u73b0\uff0c\u5e76\u9644\u4e0a\u4e86\u66f4\u591a\u5173\u4e8e\u7406\u8bba\u7684\u6ce8\u91ca\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\u56fe\u5c42\u7684\u751f\u6210\u987a\u5e8f\u4e0e\u68c0\u67e5\u70b9\u7684\u751f\u6210\u987a\u5e8f\u76f8\u540c\u3002
\n\u5b83\u5728\u56fe\u5c42\u4e0d\u53ef\u7528_^_0_^_\u65f6\u7ed9\u51fa\uff1b\u6211\u4eec\u5c06\u56fe\u5c42\u7d22\u5f15\u7528\u4f5c NeoX\uff0c\u5e76\u4e14\u5728\u5b9e\u73b0\u4e2d\u4e0d\u9700\u8981\u4e24\u4e2a\u8f6c\u6362\u5c42\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
\u590d\u5236\u7f13\u5b58\u56fe\u5c42\u6bd4\u521d\u59cb\u5316\u65b0\u56fe\u5c42\u8981\u5feb\uff0c\u56e0\u4e3a\u521d\u59cb\u5316\u53c2\u6570\u9700\u8981\u65f6\u95f4\u3002
\n\u8fd4\u56de\u521b\u5efa\u7684\u56fe\u5c42\u6216\u7f13\u5b58\u56fe\u5c42\u7684\u526f\u672c
We move the layer to the device and convert it to the correct data type
\nReturns the prepared layer
\u6211\u4eec\u5c06\u56fe\u5c42\u79fb\u52a8\u5230\u8bbe\u5907\u5e76\u5c06\u5176\u8f6c\u6362\u4e3a\u6b63\u786e\u7684\u6570\u636e\u7c7b\u578b
\n\u8fd4\u56de\u51c6\u5907\u597d\u7684\u56fe\u5c42
\n": "\n", "\n
This function implements layer transformations after loading the checkpoint.
\nCurrently, it only applies the int8 quantization.
\nReturns the prepared layer
\u6b64\u51fd\u6570\u5728\u52a0\u8f7d\u68c0\u67e5\u70b9\u540e\u5b9e\u73b0\u5c42\u8f6c\u6362\u3002
\n\u76ee\u524d\uff0c\u5b83\u4ec5\u5e94\u7528 int8 \u91cf\u5316\u3002
\n\u8fd4\u56de\u51c6\u5907\u597d\u7684\u56fe\u5c42
_^_0_^_
\n": "_^_0_^_
\n", "Code to load the checkpoint
\n": "\u52a0\u8f7d\u68c0\u67e5\u70b9\u7684\u4ee3\u7801
\n", "Readout layer
\n": "\u8bfb\u51fa\u5c42
\n", "\n": "\n", "_^_0_^_
\n": "_^_0_^_
\n", "Add RoPE embeddings
\n": "\u6dfb\u52a0\u7ef3\u7d22\u5d4c\u5165
\n", "Add head dimension
\n": "\u6dfb\u52a0\u5934\u90e8\u5c3a\u5bf8
\n", "Add them and the residual connection
\n": "\u6dfb\u52a0\u5b83\u4eec\u548c\u5269\u4f59\u7684\u8fde\u63a5
\n", "Apply mask
\n": "\u6d82\u62b9\u9762\u819c
\n", "Attention layer
\n": "\u6ce8\u610f\u5c42
\n", "Attention output transform
\n": "\u6ce8\u610f\u529b\u8f93\u51fa\u53d8\u6362
\n", "Attention query, key and value transform
\n": "\u6ce8\u610f\u529b\u67e5\u8be2\u3001\u5173\u952e\u548c\u4ef7\u503c\u8f6c\u6362
\n", "Attention scaling factor
\n": "\u6ce8\u610f\u529b\u7f29\u653e\u7cfb\u6570
\n", "Attention softmax
\n": "\u6ce8\u610f softmax
\n", "Attention softmax module
\n": "\u6ce8\u610f softmax \u6a21\u5757
\n", "Base for _^_0_^_
\n": "\u57fa\u5730_^_0_^_
\n", "Cache _^_0_^_ and _^_1_^_
\n": "\u7f13\u5b58_^_0_^_\u548c_^_1_^_
\n", "Cache them
\n": "\u7f13\u5b58\u5b83\u4eec
\n", "Calculate _^_0_^_ and _^_1_^_ in fp32
\n": "_^_1_^_\u5728 fp32 \u4e2d\u8ba1\u7b97_^_0_^_\u548c
\n", "Concatenate so that for row _^_0_^_ we have
\n_^_1_^_
\n": "\u8fde\u63a5\u8fd9\u6837_^_0_^_\u6211\u4eec\u5c31\u6709 row
\n_^_1_^_
\n", "Concatenate the past
\n": "\u4e32\u8054\u8fc7\u53bb
\n", "Concatenate with features that didn't get RoPE embeddings
\n": "\u8fde\u63a5\u672a\u83b7\u5f97 RoPe \u5d4c\u5165\u7684\u529f\u80fd
\n", "Contraction linear layer
\n": "\u6536\u7f29\u7ebf\u6027\u5c42
\n", "Convert the linear layers
\n": "\u8f6c\u6362\u7ebf\u6027\u56fe\u5c42
\n", "Convert to fp32 if the current dtype is fp16
\n": "\u5982\u679c\u5f53\u524d\u6570\u636e\u7c7b\u578b\u4e3a fp16\uff0c\u5219\u8f6c\u6362\u4e3a fp32
\n", "Create mask
\n": "\u521b\u5efa\u906e\u7f69
\n", "Disable auto-casting to fp16 for attention computation
\n": "\u7981\u7528\u81ea\u52a8\u6295\u5c04\u5230 fp16 \u4ee5\u8fdb\u884c\u6ce8\u610f\u529b\u8ba1\u7b97
\n", "Do not cast for bfloat
\n": "\u4e0d\u8981\u4e3a bfloat \u8fdb\u884c\u6295\u5c04
\n", "Embedding layer
\n": "\u5d4c\u5165\u5c42
\n", "Expansion linear layer
\n": "\u6269\u5c55\u7ebf\u6027\u5c42
\n", "FFN first transform
\n": "FFN \u9996\u6b21\u6539\u9020
\n", "FFN layer
\n": "FFN \u5c42
\n", "FFN second transform
\n": "FFN \u7b2c\u4e8c\u6b21\u53d8\u6362
\n", "Final linear layer
\n": "\u6700\u540e\u7684\u7ebf\u6027\u5c42
\n", "Final normalization layer
\n": "\u6700\u7ec8\u5f52\u4e00\u5316\u5c42
\n", "GELU activation
\n": "GELU \u6fc0\u6d3b
\n", "Get attention weighted values
\n": "\u83b7\u53d6\u6ce8\u610f\u529b\u52a0\u6743\u503c
\n", "Get causal mask
\n": "\u83b7\u5f97\u56e0\u679c\u53e3\u7f69
\n", "Get default values if not specified
\n": "\u5982\u679c\u672a\u6307\u5b9a\uff0c\u5219\u83b7\u53d6\u9ed8\u8ba4\u503c
\n", "Get position indexes _^_0_^_
\n": "\u83b7\u53d6\u5934\u5bf8\u6307\u6570_^_0_^_
\n", "Get query, key and value embeddings (all concatenated). The last dimension size will change from n_hidden -> _^_0_^_
\n": "\u83b7\u53d6\u67e5\u8be2\u3001\u952e\u548c\u503c\u5d4c\u5165\uff08\u5168\u90e8\u4e32\u8054\uff09\u3002\u6700\u540e\u4e00\u4e2a\u7ef4\u5ea6\u5927\u5c0f\u5c06\u4ece n_hidden \u66f4\u6539\u4e3a->_^_0_^_
\n", "Get the actual sequence length
\n": "\u83b7\u53d6\u5b9e\u9645\u5e8f\u5217\u957f\u5ea6
\n", "Get the past keys and values. These will have shape _^_0_^_
\n": "\u83b7\u53d6\u8fc7\u53bb\u7684\u952e\u548c\u503c\u3002\u8fd9\u4e9b\u4f1a\u6709\u5f62\u72b6_^_0_^_
\n", "Get the sin and cos values from the cache
\n": "\u4ece\u7f13\u5b58\u4e2d\u83b7\u53d6 sin \u548c cos \u503c
\n", "Get the state id's. We use to retrieve previous states and store the next states
\n": "\u83b7\u53d6\u72b6\u6001 ID\u3002\u6211\u4eec\u7528\u5b83\u6765\u68c0\u7d22\u4ee5\u524d\u7684\u72b6\u6001\u5e76\u5b58\u50a8\u4e0b\u4e00\u4e2a\u72b6\u6001
\n", "If there's cache
\n": "\u5982\u679c\u6709\u7f13\u5b58
\n", "If we are caching the states of previous tokens
\n": "\u5982\u679c\u6211\u4eec\u6b63\u5728\u7f13\u5b58\u4e4b\u524d\u4ee4\u724c\u7684\u72b6\u6001
\n", "Initialize _^_0_^_
\n": "\u521d\u59cb\u5316_^_0_^_
\n", "Initialize _^_0_^_ and _^_1_^_ cache
\n": "\u521d\u59cb\u5316_^_0_^_\u5e76_^_1_^_\u7f13\u5b58
\n", "Layer norm before FFN
\n": "FFN \u4e4b\u524d\u7684\u5206\u5c42\u89c4\u8303
\n", "Layer norm before attention
\n": "\u6ce8\u610f\u4e4b\u524d\u5148\u8fdb\u884c\u5206\u5c42\u89c4\u8303
\n", "Layer normalization before FFN
\n": "FFN \u4e4b\u524d\u7684\u5c42\u6807\u51c6\u5316
\n", "Layer normalization before attention
\n": "\u6ce8\u610f\u4e4b\u524d\u7684\u56fe\u5c42\u89c4\u8303\u5316
\n", "Linear layer for query, key and value
\n": "\u7528\u4e8e\u67e5\u8be2\u3001\u952e\u548c\u503c\u7684\u7ebf\u6027\u56fe\u5c42
\n", "NeoX runs attention and feedforward network in parallel
\n": "NeoX \u5e76\u884c\u8fd0\u884c\u6ce8\u610f\u529b\u548c\u524d\u9988\u7f51\u7edc
\n", "No cache - simply add RoPE embeddings
\n": "\u6ca1\u6709\u7f13\u5b58-\u53ea\u9700\u6dfb\u52a0 RoPe \u5d4c\u5165\u5373\u53ef
\n", "Number of features for RoPE
\n": "ROPE \u7684\u8981\u7d20\u6570\u91cf
\n", "Number of features per head
\n": "\u6bcf\u5934\u7279\u5f81\u6570
\n", "Offset of the current embeddings
\n": "\u5f53\u524d\u5d4c\u5165\u7684\u504f\u79fb\u91cf
\n", "Only convert the linear layers in the transformer layers
\n": "\u4ec5\u8f6c\u6362\u53d8\u538b\u5668\u5c42\u4e2d\u7684\u7ebf\u6027\u5c42
\n", "Otherwise, use normal attention
\n": "\u5426\u5219\uff0c\u8bf7\u6b63\u5e38\u6ce8\u610f
\n", "Query and key lengths
\n": "\u67e5\u8be2\u548c\u5bc6\u94a5\u957f\u5ea6
\n", "Readout layer
\n": "\u8bfb\u51fa\u5c42
\n", "Reshape from _^_0_^_batch_size, seq_len, n_hidden`
\n": "\u4ece_^_0_^_ batch_size\u3001seq_len\u3001n_hidden \u8fdb\u884c\u91cd\u5851 `
\n", "Residual connection
\n": "\u5269\u4f59\u8fde\u63a5
\n", "Return from cache
\n": "\u4ece\u7f13\u5b58\u4e2d\u8fd4\u56de
\n", "RoPE embedding module
\n": "\u7ef3\u7d22\u5d4c\u5165\u6a21\u5757
\n", "RoPE embeddings
\n_^_0_^_for _^_1_^_
\n": "\u7ef3\u7d22\u5d4c\u5165
\n_^_0_^_\u5bf9\u4e8e_^_1_^_
\n", "Save the current state
\n": "\u4fdd\u5b58\u5f53\u524d\u72b6\u6001
\n", "Scale attention
\n": "\u7f29\u653e\u6ce8\u610f\u529b
\n", "Skip if not using int8 quantization
\n": "\u5982\u679c\u4e0d\u4f7f\u7528 int8 \u91cf\u5316\u5219\u8df3\u8fc7
\n", "Split into heads by changing the shape to _^_0_^_
\n": "\u901a\u8fc7\u5c06\u5f62\u72b6\u6539\u4e3a\u5206\u6210\u5934\u90e8_^_0_^_
\n", "Split into query, key and value each of shape _^_0_^_
\n": "\u5206\u4e3a\u67e5\u8be2\u3001\u952e\u548c\u503c\u5404\u5f62\u72b6_^_0_^_
\n", "Split the features. We apply RoPE to only _^_0_^_ features
\n": "\u62c6\u5206\u8981\u7d20\u3002\u6211\u4eec\u4ec5\u5c06 RoPe \u5e94\u7528\u4e8e\u8981_^_0_^_\u7d20
\n", "Stack them into shape _^_0_^_
\n": "\u5c06\u5b83\u4eec\u5806\u53e0\u6210\u5f62\u72b6_^_0_^_
\n", "The output is of shape _^_0_^_
\n": "\u8f93\u51fa\u7684\u5f62\u72b6\u662f\u8fd9\u6837\u7684_^_0_^_
\n", "To cache causal mask
\n": "\u7f13\u5b58\u56e0\u679c\u63a9\u7801
\n", "To store _^_0_^_ for the features
\n": "\u4e3a\u8981\u7d20\u5b58\u50a8_^_0_^_
\n", "Transformer layer
\n": "\u53d8\u538b\u5668\u5c42
\n", "Transformer layers
\n": "\u53d8\u538b\u5668\u5c42
\n", "Use _^_0_^_ defined in utilities.
\n": "\u4f7f\u7528\u5728\u5b9e\u7528\u7a0b\u5e8f\u4e2d_^_0_^_\u5b9a\u4e49\u3002
\n", "Use flash attention
\n": "\u4f7f\u7528\u95ea\u5149\u706f\u6ce8\u610f\u529b
\n", "Out implementation doesn't include dropout.
\n": "Out \u7684\u5b9e\u73b0\u4e0d\u5305\u62ec\u9000\u51fa\u3002
\n", "