{"id":1903,"date":"2024-09-22T20:25:06","date_gmt":"2024-09-22T12:25:06","guid":{"rendered":"http:\/\/www.puzi7.com\/?p=1903"},"modified":"2026-01-26T23:11:50","modified_gmt":"2026-01-26T15:11:50","slug":"llm-%e4%b9%8bdevice_map%e5%8f%82%e6%95%b0%e8%a7%a3%e6%9e%90","status":"publish","type":"post","link":"https:\/\/www.puzi7.net\/?p=1903","title":{"rendered":"LLM \u4e4bdevice_map\u53c2\u6570\u89e3\u6790"},"content":{"rendered":"<span itemprop=\"description\">\n<p>LLM\u5927\u6a21\u578b\u5728Transformer\u505a\u63a8\u7406\u662f\uff0c\u5982\u4f55\u6839\u636eGPU\u7684\u663e\u5361\u505a\u591a\u5361\u52a0\u8f7d\u6a21\u578b\u5462\uff1f\u901a\u5e38\u4f1a\u4f7f\u7528device_map\u505a\u53c2\u6570\uff0c<strong><code>device_map<\/code>\u662f<code>AutoModel.from_pretrained()<\/code>\u00a0\u65b9\u6cd5\u4e2d\u7684\u4e00\u4e2a\u91cd\u8981\u53c2\u6570\uff0c\u5b83\u7528\u4e8e\u6307\u5b9a\u6a21\u578b\u7684\u5404\u4e2a\u90e8\u4ef6\u5e94\u52a0\u8f7d\u5230\u54ea\u4e2a\u5177\u4f53\u7684\u8ba1\u7b97\u8bbe\u5907\u4e0a\uff0c\u4ee5\u5b9e\u73b0\u8d44\u6e90\u7684\u6709\u6548\u5206\u914d\u548c\u5229\u7528<\/strong>\u3002\u8fd9\u4e2a\u53c2\u6570\u5728\u8fdb\u884c\u6a21\u578b\u5e76\u884c\u6216\u5206\u5e03\u5f0f\u8bad\u7ec3\u65f6\u7279\u522b\u6709\u7528\u3002<\/p>\n<p>device_map\u53c2\u6570\u4e3b\u8981\u6709\u5982\u4e0b4\u79cd\u9009\u9879\uff1a&#8221;auto&#8221;,&#8221;balanced&#8221;,&#8221;balanced_low_0&#8243;,&#8221;sequential&#8221;<\/p>\n<ul>\n<li>\u201cauto\u201d \u548c \u201cbalanced\u201d \u5c06\u4f1a\u5728\u6240\u6709\u7684GPU\u4e0a\u5e73\u8861\u5207\u5206\u6a21\u578b\u3002\u4e3b\u8981\u662f\u6709\u53ef\u80fd\u53d1\u73b0\u66f4\u9ad8\u6548\u7684\u5206\u914d\u7b56\u7565\u3002\u201cbalanced\u201d \u53c2\u6570\u7684\u529f\u80fd\u5219\u4fdd\u6301\u7a33\u5b9a\u3002\uff08\u4e2a\u4eba\u4e0d\u63a8\u8350\u4f7f\u7528\uff09<\/li>\n<li>\u201cbalanced_low_0\u201d \u4f1a\u5728\u9664\u4e86\u7b2c\u4e00\u4e2aGPU\u4e0a\u7684\u5176\u5b83GPU\u4e0a\u5e73\u8861\u5212\u5206\u6a21\u578b\uff0c\u5e76\u4e14\u5728\u7b2c\u4e00\u4e2a GPU \u4e0a\u5360\u636e\u8f83\u5c11\u8d44\u6e90\u3002\u8fd9\u4e2a\u9009\u9879\u7b26\u5408\u9700\u8981\u5728\u7b2c\u4e00\u4e2a GPU \u4e0a\u8fdb\u884c\u989d\u5916\u64cd\u4f5c\u7684\u9700\u6c42\uff0c\u4f8b\u5982\u9700\u8981\u5728\u7b2c\u4e00\u4e2a GPU \u6267\u884c generate \u51fd\u6570\uff08\u8fed\u4ee3\u8fc7\u7a0b\uff09\u3002\uff08\u5899\u88c2\u63a8\u8350\u4f7f\u7528\uff09<\/li>\n<li>\u201csequential\u201d \u6309\u7167GPU\u7684\u987a\u5e8f\u5206\u914d\u6a21\u578b\u5206\u7247\uff0c\u4ece GPU 0 \u5f00\u59cb\uff0c\u76f4\u5230\u6700\u540e\u7684 GPU\uff08\u90a3\u4e48\u6700\u540e\u7684 GPU \u5f80\u5f80\u4e0d\u4f1a\u88ab\u5360\u6ee1\uff0c\u548c \u201cbalanced_low_0\u201d \u7684\u533a\u522b\u5c31\u662f\u7b2c\u4e00\u4e2a\u8fd8\u662f\u6700\u540e\u4e00\u4e2a\uff0c\u4ee5\u53ca\u975e\u5747\u8861\u586b\u5145\uff09\uff0c\u4f46\u662f\u6211\u5728\u5b9e\u9645\u4f7f\u7528\u5f53\u4e2dGPU 0 \u4f1a\u76f4\u63a5\u7206\u663e\u5b58\u4e86\uff08\u76f4\u63a5\u522b\u7528\u4e86\uff09<\/li>\n<\/ul>\n<p>\u793a\u4f8b\u4ee3\u7801\uff1a<\/p>\n<div>\n<blockquote>\n<div>model = AutoModel.from_pretrained(<\/div>\n<div>\u00a0 \u00a0 path,\u00a0 # \u6a21\u578b\u672c\u5730\u8def\u5f84<\/div>\n<div>\u00a0 \u00a0 torch_dtype=torch.bfloat16,<\/div>\n<div>\u00a0 \u00a0 trust_remote_code=True,<\/div>\n<div>\u00a0 \u00a0 device_map=device_map<\/div>\n<div>).eval()<\/div>\n<\/blockquote>\n<\/div>\n<div>device_map\u53ef\u4ee5\u662f\u5b57\u7b26\u4e32 &#8220;auto&#8221;,&#8221;balanced&#8221;,&#8221;balanced_low_0&#8243;,&#8221;sequential&#8221;<\/div>\n<div>\u4e5f\u53ef\u4ee5\u624b\u52a8\u914d\u7f6e\uff0c\u6307\u5b9aGPU device<\/div>\n<div>\u00a0<\/div>\n<div><strong>\u5982\u679c\u662f\u5355\u5361\uff1adevice = &#8220;cuda:2&#8221;<\/strong><\/div>\n<blockquote>\n<div><strong>model <span class=\"token operator\">=<\/span> AutoModel<span class=\"token punctuation\">.<\/span>from_pretrained<span class=\"token punctuation\">(<\/span>model_path<span class=\"token punctuation\">,<\/span> trust_remote_code<span class=\"token operator\">=<\/span><span class=\"token boolean\">True<\/span><span class=\"token punctuation\">,<\/span> device_map<span class=\"token operator\">=<\/span>device<span class=\"token punctuation\">)<\/span><\/strong><\/div>\n<\/blockquote>\n<div>\u00a0<\/div>\n<div><strong>\u5982\u679c\u662f\u591a\u5361\uff1a\uff08\u4ee5internal-vl-40B\u4e3a\u4f8b\uff09<\/strong><\/div>\n<blockquote>\n<div>def split_model(model_name):\u00a0 \u00a0<strong>#\u6839\u636e\u5361\u7684\u6570\u91cf\u548c\u6a21\u578b\u7684\u5c42\u6570\u505a\u5747\u5206<\/strong><\/div>\n<div>\u00a0 \u00a0 device_map = {}<\/div>\n<div>\u00a0 \u00a0 world_size = torch.cuda.device_count()<\/div>\n<div>\u00a0 \u00a0 world_size = 2 #\u4f7f\u7528\u4e24\u5f20\u5361\uff0c\u5e8f\u53f7\u4ece6\u5f00\u59cb<\/div>\n<div>\u00a0 \u00a0 print(world_size)<\/div>\n<div>\u00a0 \u00a0 num_layers = {<\/div>\n<div>\u00a0 \u00a0 \u00a0 \u00a0 &#8216;InternVL2-1B&#8217;: 24, &#8216;InternVL2-2B&#8217;: 24, &#8216;InternVL2-4B&#8217;: 32, &#8216;InternVL2-8B&#8217;: 32,<\/div>\n<div>\u00a0 \u00a0 \u00a0 \u00a0 &#8216;InternVL2-26B&#8217;: 48, &#8216;InternVL2-40B&#8217;: 60, &#8216;InternVL2-Llama3-76B&#8217;: 80}[model_name]<\/div>\n<div>\u00a0 \u00a0 # Since the first GPU will be used for ViT, treat it as half a GPU.<\/div>\n<div>\u00a0 \u00a0 num_layers_per_gpu = math.ceil(num_layers \/ (world_size &#8211; 0.5))<\/div>\n<div>\u00a0 \u00a0 num_layers_per_gpu = [num_layers_per_gpu] * world_size<\/div>\n<div>\u00a0 \u00a0 num_layers_per_gpu[0] = math.ceil(num_layers_per_gpu[0] * 0.5)<\/div>\n<div>\u00a0 \u00a0 print(num_layers_per_gpu)<\/div>\n<div>\u00a0 \u00a0 layer_cnt = 0<\/div>\n<div>\u00a0 \u00a0 for i, num_layer in enumerate(num_layers_per_gpu):<\/div>\n<div>\u00a0 \u00a0 \u00a0 \u00a0 i += 6<\/div>\n<div>\u00a0 \u00a0 \u00a0 \u00a0 for j in range(num_layer):<\/div>\n<div>\u00a0 \u00a0 \u00a0 \u00a0 \u00a0 \u00a0 device_map[f&#8217;language_model.model.layers.{layer_cnt}&#8217;] = i<\/div>\n<div>\u00a0 \u00a0 \u00a0 \u00a0 \u00a0 \u00a0 # print(f&#8217;layer_cnt: {layer_cnt}&#8212;-{i}&#8217;)<\/div>\n<div>\u00a0 \u00a0 \u00a0 \u00a0 \u00a0 \u00a0 layer_cnt += 1<\/div>\n<div>\u00a0 \u00a0 print(i)\u00a0 # i=6,7<\/div>\n<div>\u00a0 \u00a0 first_gpu = i &#8211; 1 #\u628a\u5176\u4ed6\u914d\u7f6e\u90fd\u4f7f\u7528\u7b2c\u4e00\u5f20\u5361<\/div>\n<div>\u00a0 \u00a0 device_map[&#8216;vision_model&#8217;] = first_gpu<\/div>\n<div>\u00a0 \u00a0 device_map[&#8216;mlp1&#8217;] = first_gpu<\/div>\n<div>\u00a0 \u00a0 device_map[&#8216;language_model.model.tok_embeddings&#8217;] = first_gpu<\/div>\n<div>\u00a0 \u00a0 device_map[&#8216;language_model.model.embed_tokens&#8217;] = first_gpu<\/div>\n<div>\u00a0 \u00a0 device_map[&#8216;language_model.output&#8217;] = first_gpu<\/div>\n<div>\u00a0 \u00a0 device_map[&#8216;language_model.model.norm&#8217;] = first_gpu<\/div>\n<div>\u00a0 \u00a0 device_map[&#8216;language_model.lm_head&#8217;] = first_gpu<\/div>\n<div>\u00a0 \u00a0 device_map[f&#8217;language_model.model.layers.{num_layers &#8211; 1}&#8217;] = first_gpu<\/div>\n<div>\u00a0 \u00a0<\/div>\n<div>\u00a0 \u00a0 return device_map<\/div>\n<div>device = split_model(&#8220;<span style=\"color: initial;\">InternVL2-40B<\/span>&#8220;)<\/div>\n<div><strong>model <span class=\"token operator\">=<\/span> AutoModel<span class=\"token punctuation\">.<\/span>from_pretrained<span class=\"token punctuation\">(<\/span>model_path<span class=\"token punctuation\">,<\/span> trust_remote_code<span class=\"token operator\">=<\/span><span class=\"token boolean\">True<\/span><span class=\"token punctuation\">,<\/span> device_map<span class=\"token operator\">=<\/span>device<\/strong><span class=\"token punctuation\"><strong>)<\/strong><br \/><\/span><\/div>\n<\/blockquote>\n<p>\u00a0<\/p>\n\n\n\n<p>\u5982\u679c\u662f\u7528\u591a\u5361\u505a\u63a8\u7406\u52a0\u8f7d\u6a21\u578b\uff0c\u9700\u8981\u8003\u8651\u6a21\u578b\u5927\u5c0f\u4ee5\u53ca\u6bcf\u5f20\u5361\u7684\u5185\u5b58\u60c5\u51b5<\/p>\n\n\n<span>","protected":false},"excerpt":{"rendered":"<p class=\"excerpt\"><span itemprop=\"description\"><\/span><\/p>\n<p class=\"more-link-p btn-align-right\"><a class=\"green zoom-btn\" href=\"https:\/\/www.puzi7.net\/?p=1903\">Read More<\/a><\/p>\n","protected":false},"author":1,"featured_media":1976,"comment_status":"closed","ping_status":"open","sticky":false,"template":"","format":"standard","meta":{"footnotes":""},"categories":[17,25,9],"tags":[58,65,86],"class_list":["post-1903","post","type-post","status-publish","format-standard","has-post-thumbnail","hentry","category-ai","category-llm","category-9","tag-device_map","tag-gpu","tag-llm","zoom-theme-has-thumb"],"_links":{"self":[{"href":"https:\/\/www.puzi7.net\/index.php?rest_route=\/wp\/v2\/posts\/1903","targetHints":{"allow":["GET"]}}],"collection":[{"href":"https:\/\/www.puzi7.net\/index.php?rest_route=\/wp\/v2\/posts"}],"about":[{"href":"https:\/\/www.puzi7.net\/index.php?rest_route=\/wp\/v2\/types\/post"}],"author":[{"embeddable":true,"href":"https:\/\/www.puzi7.net\/index.php?rest_route=\/wp\/v2\/users\/1"}],"replies":[{"embeddable":true,"href":"https:\/\/www.puzi7.net\/index.php?rest_route=%2Fwp%2Fv2%2Fcomments&post=1903"}],"version-history":[{"count":1,"href":"https:\/\/www.puzi7.net\/index.php?rest_route=\/wp\/v2\/posts\/1903\/revisions"}],"predecessor-version":[{"id":1977,"href":"https:\/\/www.puzi7.net\/index.php?rest_route=\/wp\/v2\/posts\/1903\/revisions\/1977"}],"wp:featuredmedia":[{"embeddable":true,"href":"https:\/\/www.puzi7.net\/index.php?rest_route=\/wp\/v2\/media\/1976"}],"wp:attachment":[{"href":"https:\/\/www.puzi7.net\/index.php?rest_route=%2Fwp%2Fv2%2Fmedia&parent=1903"}],"wp:term":[{"taxonomy":"category","embeddable":true,"href":"https:\/\/www.puzi7.net\/index.php?rest_route=%2Fwp%2Fv2%2Fcategories&post=1903"},{"taxonomy":"post_tag","embeddable":true,"href":"https:\/\/www.puzi7.net\/index.php?rest_route=%2Fwp%2Fv2%2Ftags&post=1903"}],"curies":[{"name":"wp","href":"https:\/\/api.w.org\/{rel}","templated":true}]}}