import{s as P,n as aa,o as ta}from"../chunks/scheduler.25b97de1.js";import{S as la,i as sa,g as i,s as e,r as w,A as ea,h as p,f as l,c as n,j as L,u as J,x as g,k as D,y as na,a as s,v as f,d as T,t as W,w as Z}from"../chunks/index.d9030fc9.js";import{C as N}from"../chunks/CodeBlock.e6cd0d95.js";import{H as O,E as oa}from"../chunks/EditOnGithub.91d95064.js";function ia(A){let o,R,X,G,M,k,m,x='<a href="https://github.com/pytorch/ao" rel="nofollow">TorchAO</a> is an architecture optimization library for PyTorch, it provides high performance dtypes, optimization techniques and kernels for inference and training, featuring composability with native PyTorch features like <code>torch.compile</code>, FSDP etc.. Some benchmark numbers can be found <a href="https://github.com/pytorch/ao/tree/main/torchao/quantization#benchmarks" rel="nofollow">here</a>.',V,c,S="Before you begin, make sure the following libraries are installed with their latest version:",B,r,C,d,H="By default, the weights are loaded in full precision (torch.float32) regardless of the actual data type the weights are stored in such as torch.float16. Set <code>torch_dtype=&quot;auto&quot;</code> to load the weights in the data type defined in a model’s <code>config.json</code> file to automatically load the most memory-optimal data type.",Y,h,z,u,I,y,$='torchao quantization is implemented with <a href="https://pytorch.org/docs/stable/notes/extending.html#subclassing-torch-tensor" rel="nofollow">tensor subclasses</a>, it only work with huggingface non-safetensor serialization and deserialization. It relies on <code>torch.load(..., weights_only=True)</code> to avoid arbitrary user code execution during load time and use <a href="https://pytorch.org/docs/stable/notes/serialization.html#torch.serialization.add_safe_globals" rel="nofollow">add_safe_globals</a> to allowlist some known user functions.',F,b,K="The reason why it does not support safe tensor serialization is that wrapper tensor subclass allows maximum flexibility so we want to make sure the effort of supporting new format of quantized Tensor is low, while safe tensor optimizes for maximum safety (no user code execution), it also means we have to make sure to manually support new quantization format.",v,j,Q,U,q,_,E;return M=new O({props:{title:"TorchAO",local:"torchao",headingTag:"h1"}}),r=new N({props:{code:"JTIzJTIwVXBkYXRpbmclMjAlRjAlOUYlQTQlOTclMjBUcmFuc2Zvcm1lcnMlMjB0byUyMHRoZSUyMGxhdGVzdCUyMHZlcnNpb24lMkMlMjBhcyUyMHRoZSUyMGV4YW1wbGUlMjBzY3JpcHQlMjBiZWxvdyUyMHVzZXMlMjB0aGUlMjBuZXclMjBhdXRvJTIwY29tcGlsYXRpb24lMEFwaXAlMjBpbnN0YWxsJTIwLS11cGdyYWRlJTIwdG9yY2glMjB0b3JjaGFvJTIwdHJhbnNmb3JtZXJz",highlighted:`<span class="hljs-comment"># Updating 🤗 Transformers to the latest version, as the example script below uses the new auto compilation</span>
pip install --upgrade torch torchao transformers`,wrap:!1}}),h=new N({props:{code:"aW1wb3J0JTIwdG9yY2glMEFmcm9tJTIwdHJhbnNmb3JtZXJzJTIwaW1wb3J0JTIwVG9yY2hBb0NvbmZpZyUyQyUyMEF1dG9Nb2RlbEZvckNhdXNhbExNJTJDJTIwQXV0b1Rva2VuaXplciUwQSUwQW1vZGVsX25hbWUlMjAlM0QlMjAlMjJtZXRhLWxsYW1hJTJGTWV0YS1MbGFtYS0zLThCJTIyJTBBJTIzJTIwV2UlMjBzdXBwb3J0JTIwaW50NF93ZWlnaHRfb25seSUyQyUyMGludDhfd2VpZ2h0X29ubHklMjBhbmQlMjBpbnQ4X2R5bmFtaWNfYWN0aXZhdGlvbl9pbnQ4X3dlaWdodCUwQSUyMyUyME1vcmUlMjBleGFtcGxlcyUyMGFuZCUyMGRvY3VtZW50YXRpb25zJTIwZm9yJTIwYXJndW1lbnRzJTIwY2FuJTIwYmUlMjBmb3VuZCUyMGluJTIwaHR0cHMlM0ElMkYlMkZnaXRodWIuY29tJTJGcHl0b3JjaCUyRmFvJTJGdHJlZSUyRm1haW4lMkZ0b3JjaGFvJTJGcXVhbnRpemF0aW9uJTIzb3RoZXItYXZhaWxhYmxlLXF1YW50aXphdGlvbi10ZWNobmlxdWVzJTBBcXVhbnRpemF0aW9uX2NvbmZpZyUyMCUzRCUyMFRvcmNoQW9Db25maWcoJTIyaW50NF93ZWlnaHRfb25seSUyMiUyQyUyMGdyb3VwX3NpemUlM0QxMjgpJTBBcXVhbnRpemVkX21vZGVsJTIwJTNEJTIwQXV0b01vZGVsRm9yQ2F1c2FsTE0uZnJvbV9wcmV0cmFpbmVkKG1vZGVsX25hbWUlMkMlMjB0b3JjaF9kdHlwZSUzRCUyMmF1dG8lMjIlMkMlMjBkZXZpY2VfbWFwJTNEJTIyYXV0byUyMiUyQyUyMHF1YW50aXphdGlvbl9jb25maWclM0RxdWFudGl6YXRpb25fY29uZmlnKSUwQSUwQXRva2VuaXplciUyMCUzRCUyMEF1dG9Ub2tlbml6ZXIuZnJvbV9wcmV0cmFpbmVkKG1vZGVsX25hbWUpJTBBaW5wdXRfdGV4dCUyMCUzRCUyMCUyMldoYXQlMjBhcmUlMjB3ZSUyMGhhdmluZyUyMGZvciUyMGRpbm5lciUzRiUyMiUwQWlucHV0X2lkcyUyMCUzRCUyMHRva2VuaXplcihpbnB1dF90ZXh0JTJDJTIwcmV0dXJuX3RlbnNvcnMlM0QlMjJwdCUyMikudG8oJTIyY3VkYSUyMiklMEElMEElMjMlMjBhdXRvLWNvbXBpbGUlMjB0aGUlMjBxdWFudGl6ZWQlMjBtb2RlbCUyMHdpdGglMjAlNjBjYWNoZV9pbXBsZW1lbnRhdGlvbiUzRCUyMnN0YXRpYyUyMiU2MCUyMHRvJTIwZ2V0JTIwc3BlZWR1cCUwQW91dHB1dCUyMCUzRCUyMHF1YW50aXplZF9tb2RlbC5nZW5lcmF0ZSgqKmlucHV0X2lkcyUyQyUyMG1heF9uZXdfdG9rZW5zJTNEMTAlMkMlMjBjYWNoZV9pbXBsZW1lbnRhdGlvbiUzRCUyMnN0YXRpYyUyMiklMEFwcmludCh0b2tlbml6ZXIuZGVjb2RlKG91dHB1dCU1QjAlNUQlMkMlMjBza2lwX3NwZWNpYWxfdG9rZW5zJTNEVHJ1ZSkpJTBBJTBBJTIzJTIwYmVuY2htYXJrJTIwdGhlJTIwcGVyZm9ybWFuY2UlMEFpbXBvcnQlMjB0b3JjaC51dGlscy5iZW5jaG1hcmslMjBhcyUyMGJlbmNobWFyayUwQSUwQWRlZiUyMGJlbmNobWFya19mbihmJTJDJTIwKmFyZ3MlMkMlMjAqKmt3YXJncyklM0ElMEElMjAlMjAlMjAlMjAlMjMlMjBNYW51YWwlMjB3YXJtdXAlMEElMjAlMjAlMjAlMjBmb3IlMjBfJTIwaW4lMjByYW5nZSg1KSUzQSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGYoKmFyZ3MlMkMlMjAqKmt3YXJncyklMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMEElMjAlMjAlMjAlMjB0MCUyMCUzRCUyMGJlbmNobWFyay5UaW1lciglMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBzdG10JTNEJTIyZigqYXJncyUyQyUyMCoqa3dhcmdzKSUyMiUyQyUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGdsb2JhbHMlM0QlN0IlMjJhcmdzJTIyJTNBJTIwYXJncyUyQyUyMCUyMmt3YXJncyUyMiUzQSUyMGt3YXJncyUyQyUyMCUyMmYlMjIlM0ElMjBmJTdEJTJDJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIwJTIwbnVtX3RocmVhZHMlM0R0b3JjaC5nZXRfbnVtX3RocmVhZHMoKSUyQyUwQSUyMCUyMCUyMCUyMCklMEElMjAlMjAlMjAlMjByZXR1cm4lMjBmJTIyJTdCKHQwLmJsb2NrZWRfYXV0b3JhbmdlKCkubWVhbiklM0EuM2YlN0QlMjIlMEElMEFNQVhfTkVXX1RPS0VOUyUyMCUzRCUyMDEwMDAlMEFwcmludCglMjJpbnQ0d28tMTI4JTIwbW9kZWwlM0ElMjIlMkMlMjBiZW5jaG1hcmtfZm4ocXVhbnRpemVkX21vZGVsLmdlbmVyYXRlJTJDJTIwKippbnB1dF9pZHMlMkMlMjBtYXhfbmV3X3Rva2VucyUzRE1BWF9ORVdfVE9LRU5TJTJDJTIwY2FjaGVfaW1wbGVtZW50YXRpb24lM0QlMjJzdGF0aWMlMjIpKSUwQSUwQWJmMTZfbW9kZWwlMjAlM0QlMjBBdXRvTW9kZWxGb3JDYXVzYWxMTS5mcm9tX3ByZXRyYWluZWQobW9kZWxfbmFtZSUyQyUyMGRldmljZV9tYXAlM0QlMjJjdWRhJTIyJTJDJTIwdG9yY2hfZHR5cGUlM0R0b3JjaC5iZmxvYXQxNiklMEFvdXRwdXQlMjAlM0QlMjBiZjE2X21vZGVsLmdlbmVyYXRlKCoqaW5wdXRfaWRzJTJDJTIwbWF4X25ld190b2tlbnMlM0QxMCUyQyUyMGNhY2hlX2ltcGxlbWVudGF0aW9uJTNEJTIyc3RhdGljJTIyKSUyMCUyMyUyMGF1dG8tY29tcGlsZSUwQXByaW50KCUyMmJmMTYlMjBtb2RlbCUzQSUyMiUyQyUyMGJlbmNobWFya19mbihiZjE2X21vZGVsLmdlbmVyYXRlJTJDJTIwKippbnB1dF9pZHMlMkMlMjBtYXhfbmV3X3Rva2VucyUzRE1BWF9ORVdfVE9LRU5TJTJDJTIwY2FjaGVfaW1wbGVtZW50YXRpb24lM0QlMjJzdGF0aWMlMjIpKSUwQQ==",highlighted:`<span class="hljs-keyword">import</span> torch
<span class="hljs-keyword">from</span> transformers <span class="hljs-keyword">import</span> TorchAoConfig, AutoModelForCausalLM, AutoTokenizer

model_name = <span class="hljs-string">&quot;meta-llama/Meta-Llama-3-8B&quot;</span>
<span class="hljs-comment"># We support int4_weight_only, int8_weight_only and int8_dynamic_activation_int8_weight</span>
<span class="hljs-comment"># More examples and documentations for arguments can be found in https://github.com/pytorch/ao/tree/main/torchao/quantization#other-available-quantization-techniques</span>
quantization_config = TorchAoConfig(<span class="hljs-string">&quot;int4_weight_only&quot;</span>, group_size=<span class="hljs-number">128</span>)
quantized_model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=<span class="hljs-string">&quot;auto&quot;</span>, device_map=<span class="hljs-string">&quot;auto&quot;</span>, quantization_config=quantization_config)

tokenizer = AutoTokenizer.from_pretrained(model_name)
input_text = <span class="hljs-string">&quot;What are we having for dinner?&quot;</span>
input_ids = tokenizer(input_text, return_tensors=<span class="hljs-string">&quot;pt&quot;</span>).to(<span class="hljs-string">&quot;cuda&quot;</span>)

<span class="hljs-comment"># auto-compile the quantized model with \`cache_implementation=&quot;static&quot;\` to get speedup</span>
output = quantized_model.generate(**input_ids, max_new_tokens=<span class="hljs-number">10</span>, cache_implementation=<span class="hljs-string">&quot;static&quot;</span>)
<span class="hljs-built_in">print</span>(tokenizer.decode(output[<span class="hljs-number">0</span>], skip_special_tokens=<span class="hljs-literal">True</span>))

<span class="hljs-comment"># benchmark the performance</span>
<span class="hljs-keyword">import</span> torch.utils.benchmark <span class="hljs-keyword">as</span> benchmark

<span class="hljs-keyword">def</span> <span class="hljs-title function_">benchmark_fn</span>(<span class="hljs-params">f, *args, **kwargs</span>):
    <span class="hljs-comment"># Manual warmup</span>
    <span class="hljs-keyword">for</span> _ <span class="hljs-keyword">in</span> <span class="hljs-built_in">range</span>(<span class="hljs-number">5</span>):
        f(*args, **kwargs)
        
    t0 = benchmark.Timer(
        stmt=<span class="hljs-string">&quot;f(*args, **kwargs)&quot;</span>,
        <span class="hljs-built_in">globals</span>={<span class="hljs-string">&quot;args&quot;</span>: args, <span class="hljs-string">&quot;kwargs&quot;</span>: kwargs, <span class="hljs-string">&quot;f&quot;</span>: f},
        num_threads=torch.get_num_threads(),
    )
    <span class="hljs-keyword">return</span> <span class="hljs-string">f&quot;<span class="hljs-subst">{(t0.blocked_autorange().mean):<span class="hljs-number">.3</span>f}</span>&quot;</span>

MAX_NEW_TOKENS = <span class="hljs-number">1000</span>
<span class="hljs-built_in">print</span>(<span class="hljs-string">&quot;int4wo-128 model:&quot;</span>, benchmark_fn(quantized_model.generate, **input_ids, max_new_tokens=MAX_NEW_TOKENS, cache_implementation=<span class="hljs-string">&quot;static&quot;</span>))

bf16_model = AutoModelForCausalLM.from_pretrained(model_name, device_map=<span class="hljs-string">&quot;cuda&quot;</span>, torch_dtype=torch.bfloat16)
output = bf16_model.generate(**input_ids, max_new_tokens=<span class="hljs-number">10</span>, cache_implementation=<span class="hljs-string">&quot;static&quot;</span>) <span class="hljs-comment"># auto-compile</span>
<span class="hljs-built_in">print</span>(<span class="hljs-string">&quot;bf16 model:&quot;</span>, benchmark_fn(bf16_model.generate, **input_ids, max_new_tokens=MAX_NEW_TOKENS, cache_implementation=<span class="hljs-string">&quot;static&quot;</span>))
`,wrap:!1}}),u=new O({props:{title:"Serialization and Deserialization",local:"serialization-and-deserialization",headingTag:"h2"}}),j=new N({props:{code:"JTIzJTIwc2F2ZSUyMHF1YW50aXplZCUyMG1vZGVsJTIwbG9jYWxseSUwQW91dHB1dF9kaXIlMjAlM0QlMjAlMjJsbGFtYTMtOGItaW50NHdvLTEyOCUyMiUwQXF1YW50aXplZF9tb2RlbC5zYXZlX3ByZXRyYWluZWQob3V0cHV0X2RpciUyQyUyMHNhZmVfc2VyaWFsaXphdGlvbiUzREZhbHNlKSUwQSUwQSUyMyUyMHB1c2glMjB0byUyMGh1Z2dpbmdmYWNlJTIwaHViJTBBJTIzJTIwc2F2ZV90byUyMCUzRCUyMCUyMiU3QnVzZXJfaWQlN0QlMkZsbGFtYTMtOGItaW50NHdvLTEyOCUyMiUwQSUyMyUyMHF1YW50aXplZF9tb2RlbC5wdXNoX3RvX2h1YihzYXZlX3RvJTJDJTIwc2FmZV9zZXJpYWxpemF0aW9uJTNERmFsc2UpJTBBJTBBJTIzJTIwbG9hZCUyMHF1YW50aXplZCUyMG1vZGVsJTBBY2twdF9pZCUyMCUzRCUyMCUyMmxsYW1hMy04Yi1pbnQ0d28tMTI4JTIyJTIwJTIwJTIzJTIwb3IlMjBodWdnaW5nZmFjZSUyMGh1YiUyMG1vZGVsJTIwaWQlMEFsb2FkZWRfcXVhbnRpemVkX21vZGVsJTIwJTNEJTIwQXV0b01vZGVsRm9yQ2F1c2FsTE0uZnJvbV9wcmV0cmFpbmVkKGNrcHRfaWQlMkMlMjBkZXZpY2VfbWFwJTNEJTIyY3VkYSUyMiklMEElMEElMEElMjMlMjBjb25maXJtJTIwdGhlJTIwc3BlZWR1cCUwQWxvYWRlZF9xdWFudGl6ZWRfbW9kZWwlMjAlM0QlMjB0b3JjaC5jb21waWxlKGxvYWRlZF9xdWFudGl6ZWRfbW9kZWwlMkMlMjBtb2RlJTNEJTIybWF4LWF1dG90dW5lJTIyKSUwQXByaW50KCUyMmxvYWRlZCUyMGludDR3by0xMjglMjBtb2RlbCUzQSUyMiUyQyUyMGJlbmNobWFya19mbihsb2FkZWRfcXVhbnRpemVkX21vZGVsLmdlbmVyYXRlJTJDJTIwKippbnB1dF9pZHMlMkMlMjBtYXhfbmV3X3Rva2VucyUzRE1BWF9ORVdfVE9LRU5TKSk=",highlighted:`<span class="hljs-comment"># save quantized model locally</span>
output_dir = <span class="hljs-string">&quot;llama3-8b-int4wo-128&quot;</span>
quantized_model.save_pretrained(output_dir, safe_serialization=<span class="hljs-literal">False</span>)

<span class="hljs-comment"># push to huggingface hub</span>
<span class="hljs-comment"># save_to = &quot;{user_id}/llama3-8b-int4wo-128&quot;</span>
<span class="hljs-comment"># quantized_model.push_to_hub(save_to, safe_serialization=False)</span>

<span class="hljs-comment"># load quantized model</span>
ckpt_id = <span class="hljs-string">&quot;llama3-8b-int4wo-128&quot;</span>  <span class="hljs-comment"># or huggingface hub model id</span>
loaded_quantized_model = AutoModelForCausalLM.from_pretrained(ckpt_id, device_map=<span class="hljs-string">&quot;cuda&quot;</span>)


<span class="hljs-comment"># confirm the speedup</span>
loaded_quantized_model = torch.<span class="hljs-built_in">compile</span>(loaded_quantized_model, mode=<span class="hljs-string">&quot;max-autotune&quot;</span>)
<span class="hljs-built_in">print</span>(<span class="hljs-string">&quot;loaded int4wo-128 model:&quot;</span>, benchmark_fn(loaded_quantized_model.generate, **input_ids, max_new_tokens=MAX_NEW_TOKENS))`,wrap:!1}}),U=new oa({props:{source:"https://github.com/huggingface/transformers/blob/main/docs/source/en/quantization/torchao.md"}}),{c(){o=i("meta"),R=e(),X=i("p"),G=e(),w(M.$$.fragment),k=e(),m=i("p"),m.innerHTML=x,V=e(),c=i("p"),c.textContent=S,B=e(),w(r.$$.fragment),C=e(),d=i("p"),d.innerHTML=H,Y=e(),w(h.$$.fragment),z=e(),w(u.$$.fragment),I=e(),y=i("p"),y.innerHTML=$,F=e(),b=i("p"),b.textContent=K,v=e(),w(j.$$.fragment),Q=e(),w(U.$$.fragment),q=e(),_=i("p"),this.h()},l(a){const t=ea("svelte-u9bgzb",document.head);o=p(t,"META",{name:!0,content:!0}),t.forEach(l),R=n(a),X=p(a,"P",{}),L(X).forEach(l),G=n(a),J(M.$$.fragment,a),k=n(a),m=p(a,"P",{"data-svelte-h":!0}),g(m)!=="svelte-1w5ww4g"&&(m.innerHTML=x),V=n(a),c=p(a,"P",{"data-svelte-h":!0}),g(c)!=="svelte-1wrsrbm"&&(c.textContent=S),B=n(a),J(r.$$.fragment,a),C=n(a),d=p(a,"P",{"data-svelte-h":!0}),g(d)!=="svelte-9fdmcv"&&(d.innerHTML=H),Y=n(a),J(h.$$.fragment,a),z=n(a),J(u.$$.fragment,a),I=n(a),y=p(a,"P",{"data-svelte-h":!0}),g(y)!=="svelte-14dwjgd"&&(y.innerHTML=$),F=n(a),b=p(a,"P",{"data-svelte-h":!0}),g(b)!=="svelte-10wmnl9"&&(b.textContent=K),v=n(a),J(j.$$.fragment,a),Q=n(a),J(U.$$.fragment,a),q=n(a),_=p(a,"P",{}),L(_).forEach(l),this.h()},h(){D(o,"name","hf:doc:metadata"),D(o,"content",pa)},m(a,t){na(document.head,o),s(a,R,t),s(a,X,t),s(a,G,t),f(M,a,t),s(a,k,t),s(a,m,t),s(a,V,t),s(a,c,t),s(a,B,t),f(r,a,t),s(a,C,t),s(a,d,t),s(a,Y,t),f(h,a,t),s(a,z,t),f(u,a,t),s(a,I,t),s(a,y,t),s(a,F,t),s(a,b,t),s(a,v,t),f(j,a,t),s(a,Q,t),f(U,a,t),s(a,q,t),s(a,_,t),E=!0},p:aa,i(a){E||(T(M.$$.fragment,a),T(r.$$.fragment,a),T(h.$$.fragment,a),T(u.$$.fragment,a),T(j.$$.fragment,a),T(U.$$.fragment,a),E=!0)},o(a){W(M.$$.fragment,a),W(r.$$.fragment,a),W(h.$$.fragment,a),W(u.$$.fragment,a),W(j.$$.fragment,a),W(U.$$.fragment,a),E=!1},d(a){a&&(l(R),l(X),l(G),l(k),l(m),l(V),l(c),l(B),l(C),l(d),l(Y),l(z),l(I),l(y),l(F),l(b),l(v),l(Q),l(q),l(_)),l(o),Z(M,a),Z(r,a),Z(h,a),Z(u,a),Z(j,a),Z(U,a)}}}const pa='{"title":"TorchAO","local":"torchao","sections":[{"title":"Serialization and Deserialization","local":"serialization-and-deserialization","sections":[],"depth":2}],"depth":1}';function Ma(A){return ta(()=>{new URLSearchParams(window.location.search).get("fw")}),[]}class ha extends la{constructor(o){super(),sa(this,o,Ma,ia,P,{})}}export{ha as component};
