accelerate/pr_4103/en/_app/immutable/nodes/20.CSed2Ucy.js

import"../chunks/DsnmJJEf.js";import{i as I,h as V,C as Z,H as i,a as l,E as W,s as A}from"../chunks/ClLAY5C0.js";import{p as X,o as k,s as e,f as R,a as n,b as B,c as s,n as F}from"../chunks/BhevU81Y.js";import{T as r}from"../chunks/L24edx8-.js";const S='{"title":"Training on TPUs","local":"training-on-tpus","sections":[{"title":"Training in a Notebook","local":"training-in-a-notebook","sections":[],"depth":2},{"title":"Mixed Precision and Global Variables","local":"mixed-precision-and-global-variables","sections":[],"depth":2},{"title":"Training Times on TPUs","local":"training-times-on-tpus","sections":[],"depth":2}],"depth":1}';var z=s('<meta name="hf:doc:metadata"/>'),Q=s(`<p>This code snippet is based off the one from the <code>simple_nlp_example</code> notebook found <a href="https://github.com/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_nlp_example.ipynb" rel="nofollow">here</a> with slight
modifications for the sake of simplicity</p>`),x=s("<p>The <code>notebook_launcher</code> will default to 8 processes if Accelerate has been configured for a TPU</p>"),N=s(`<p>The above workaround is only needed when launching a TPU instance from a Jupyter Notebook on a low-resource server such as Google Colaboratory or Kaggle. If
using a script or launching on a much beefier server declaring the model beforehand is not needed.</p>`),E=s("<p>Just because the memory is allocated does not mean it will be used or that the batch size will increase when going back to your training dataloader.</p>"),H=s(`<p></p> <!> <!> <p>Training on TPUs can be slightly different from training on multi-gpu, even with Accelerate. This guide aims to show you
where you should be careful and why, as well as the best practices in general.</p> <!> <p>The main carepoint when training on TPUs comes from the <a href="/docs/accelerate/pr_4103/en/package_reference/launchers#accelerate.notebook_launcher">notebook_launcher()</a>. As mentioned in the <a href="../usage_guides/notebook">notebook tutorial</a>, you need to
restructure your training code into a function that can get passed to the <a href="/docs/accelerate/pr_4103/en/package_reference/launchers#accelerate.notebook_launcher">notebook_launcher()</a> function and be careful about not declaring any tensors on the GPU.</p> <p>While on a TPU that last part is not as important, a critical part to understand is that when you launch code from a notebook you do so through a process called <strong>forking</strong>.
When launching from the command-line, you perform <strong>spawning</strong>, where a python process is not currently running and you <em>spawn</em> a new process in. Since your Jupyter notebook is already
utilizing a python process, you need to <em>fork</em> a new process from it to launch your code.</p> <p>Where this becomes important is in regard to declaring your model. On forked TPU processes, it is recommended that you instantiate your model <em>once</em> and pass this into your
training function. This is different than training on GPUs where you create <code>n</code> models that have their gradients synced and back-propagated at certain moments. Instead, one
model instance is shared between all the nodes and it is passed back and forth. This is important especially when training on low-resource TPUs such as those provided in Kaggle kernels or
on Google Colaboratory.</p> <p>Below is an example of a training function passed to the <a href="/docs/accelerate/pr_4103/en/package_reference/launchers#accelerate.notebook_launcher">notebook_launcher()</a> if training on CPUs or GPUs:</p> <!> <!> <!> <!> <p>If you use this example and declare the model <em>inside</em> the training loop, then on a low-resource system you will potentially see an error
like:</p> <!> <p>This error is <em>extremely</em> cryptic but the basic explanation is you ran out of system RAM. You can avoid this entirely by reconfiguring the training function to
accept a single <code>model</code> argument, and declare it in an outside cell:</p> <!> <!> <p>And finally calling the training function with:</p> <!> <!> <!> <p>As mentioned in the <a href="../usage_guides/mixed_precision">mixed precision tutorial</a>, Accelerate supports fp16 and bf16, both of which can be used on TPUs.
That being said, ideally <code>bf16</code> should be utilized as it is extremely efficient to use.</p> <p>There are two “layers” when using <code>bf16</code> and Accelerate on TPUs, at the base level and at the operation level.</p> <p>At the base level, this is enabled when passing <code>mixed_precision="bf16"</code> to <code>Accelerator</code>, such as:</p> <!> <p>By default, this will cast <code>torch.float</code> and <code>torch.double</code> to <code>bfloat16</code> on TPUs.
The specific configuration being set is an environmental variable of <code>XLA_USE_BF16</code> is set to <code>1</code>.</p> <p>There is a further configuration you can perform which is setting the <code>XLA_DOWNCAST_BF16</code> environmental variable. If set to <code>1</code>, then <code>torch.float</code> is <code>bfloat16</code> and <code>torch.double</code> is <code>float32</code>.</p> <p>This is performed in the <code>Accelerator</code> object when passing <code>downcast_bf16=True</code>:</p> <!> <p>Using downcasting instead of bf16 everywhere is good for when you are trying to calculate metrics, log values, and more where raw bf16 tensors would be unusable.</p> <!> <p>As you launch your script, you may notice that training seems exceptionally slow at first. This is because TPUs
first run through a few batches of data to see how much memory to allocate before finally utilizing this configured
memory allocation extremely efficiently.</p> <p>If you notice that your evaluation code to calculate the metrics of your model takes longer due to a larger batch size being used,
it is recommended to keep the batch size the same as the training data if it is too slow. Otherwise the memory will reallocate to this
new batch size after the first few iterations.</p> <!> <!> <p></p>`,1);function D(Y,_){X(_,!1),k(()=>{new URLSearchParams(window.location.search).get("fw")}),I();var c=H();V("1ulx9jk",a=>{var t=z();A(t,"content",S),n(a,t)});var d=e(R(c),2);Z(d,{containerStyle:"float: right; margin-left: 10px; display: inline-flex; position: relative; z-index: 10;"});var h=e(d,2);i(h,{title:"Training on TPUs",local:"training-on-tpus",headingTag:"h1"});var y=e(h,4);i(y,{title:"Training in a Notebook",local:"training-in-a-notebook",headingTag:"h2"});var p=e(y,10);r(p,{children:(a,t)=>{var o=Q();n(a,o)},$$slots:{default:!0}});var M=e(p,2);l(M,{code:"ZGVmJTIwdHJhaW5pbmdfZnVuY3Rpb24oKSUzQSUwQSUyMCUyMCUyMCUyMCUyMyUyMEluaXRpYWxpemUlMjBhY2NlbGVyYXRvciUwQSUyMCUyMCUyMCUyMGFjY2VsZXJhdG9yJTIwJTNEJTIwQWNjZWxlcmF0b3IoKSUwQSUyMCUyMCUyMCUyMG1vZGVsJTIwJTNEJTIwQXV0b01vZGVsRm9yU2VxdWVuY2VDbGFzc2lmaWNhdGlvbi5mcm9tX3ByZXRyYWluZWQoJTIyYmVydC1iYXNlLWNhc2VkJTIyJTJDJTIwbnVtX2xhYmVscyUzRDIpJTBBJTIwJTIwJTIwJTIwdHJhaW5fZGF0YWxvYWRlciUyQyUyMGV2YWxfZGF0YWxvYWRlciUyMCUzRCUyMGNyZWF0ZV9kYXRhbG9hZGVycyglMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjB0cmFpbl9iYXRjaF9zaXplJTNEaHlwZXJwYXJhbWV0ZXJzJTVCJTIydHJhaW5fYmF0Y2hfc2l6ZSUyMiU1RCUyQyUyMGV2YWxfYmF0Y2hfc2l6ZSUzRGh5cGVycGFyYW1ldGVycyU1QiUyMmV2YWxfYmF0Y2hfc2l6ZSUyMiU1RCUwQSUyMCUyMCUyMCUyMCklMEElMEElMjAlMjAlMjAlMjAlMjMlMjBJbnN0YW50aWF0ZSUyMG9wdGltaXplciUwQSUyMCUyMCUyMCUyMG9wdGltaXplciUyMCUzRCUyMEFkYW1XKHBhcmFtcyUzRG1vZGVsLnBhcmFtZXRlcnMoKSUyQyUyMGxyJTNEaHlwZXJwYXJhbWV0ZXJzJTVCJTIybGVhcm5pbmdfcmF0ZSUyMiU1RCklMEElMEElMjAlMjAlMjAlMjAlMjMlMjBQcmVwYXJlJTIwZXZlcnl0aGluZyUwQSUyMCUyMCUyMCUyMCUyMyUyMFRoZXJlJTIwaXMlMjBubyUyMHNwZWNpZmljJTIwb3JkZXIlMjB0byUyMHJlbWVtYmVyJTJDJTIwd2UlMjBqdXN0JTIwbmVlZCUyMHRvJTIwdW5wYWNrJTIwdGhlJTIwb2JqZWN0cyUyMGluJTIwdGhlJTIwc2FtZSUyMG9yZGVyJTIwd2UlMjBnYXZlJTIwdGhlbSUyMHRvJTIwdGhlJTBBJTIwJTIwJTIwJTIwJTIzJTIwcHJlcGFyZSUyMG1ldGhvZC4lMEElMjAlMjAlMjAlMjBtb2RlbCUyQyUyMG9wdGltaXplciUyQyUyMHRyYWluX2RhdGFsb2FkZXIlMkMlMjBldmFsX2RhdGFsb2FkZXIlMjAlM0QlMjBhY2NlbGVyYXRvci5wcmVwYXJlKCUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMG1vZGVsJTJDJTIwb3B0aW1pemVyJTJDJTIwdHJhaW5fZGF0YWxvYWRlciUyQyUyMGV2YWxfZGF0YWxvYWRlciUwQSUyMCUyMCUyMCUyMCklMEElMEElMjAlMjAlMjAlMjBudW1fZXBvY2hzJTIwJTNEJTIwaHlwZXJwYXJhbWV0ZXJzJTVCJTIybnVtX2Vwb2NocyUyMiU1RCUwQSUyMCUyMCUyMCUyMCUyMyUyME5vdyUyMHdlJTIwdHJhaW4lMjB0aGUlMjBtb2RlbCUwQSUyMCUyMCUyMCUyMGZvciUyMGVwb2NoJTIwaW4lMjByYW5nZShudW1fZXBvY2hzKSUzQSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMG1vZGVsLnRyYWluKCklMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBmb3IlMjBzdGVwJTJDJTIwYmF0Y2glMjBpbiUyMGVudW1lcmF0ZSh0cmFpbl9kYXRhbG9hZGVyKSUzQSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMG91dHB1dHMlMjAlM0QlMjBtb2RlbCgqKmJhdGNoKSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMGxvc3MlMjAlM0QlMjBvdXRwdXRzLmxvc3MlMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBhY2NlbGVyYXRvci5iYWNrd2FyZChsb3NzKSUwQSUwQSUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMCUyMG9wdGltaXplci5zdGVwKCklMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjBvcHRpbWl6ZXIuemVyb19ncmFkKCk=",highlighted:`<span class="hljs-keyword">def</span> <span class="hljs-title function_">training_function</span>():
    <span class="hljs-comment"># Initialize accelerator</span>
    accelerator = Accelerator()
    model = AutoModelForSequenceClassification.from_pretrained(<span class="hljs-string">&quot;bert-base-cased&quot;</span>, num_labels=<span class="hljs-number">2</span>)
    train_dataloader, eval_dataloader = create_dataloaders(
        train_batch_size=hyperparameters[<span class="hljs-string">&quot;train_batch_size&quot;</span>], eval_batch_size=hyperparameters[<span class="hljs-string">&quot;eval_batch_size&quot;</span>]
    )

    <span class="hljs-comment"># Instantiate optimizer</span>
    optimizer = AdamW(params=model.parameters(), lr=hyperparameters[<span class="hljs-string">&quot;learning_rate&quot;</span>])

    <span class="hljs-comment"># Prepare everything</span>
    <span class="hljs-comment"># There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the</span>
    <span class="hljs-comment"># prepare method.</span>
    model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader
    )

    num_epochs = hyperparameters[<span class="hljs-string">&quot;num_epochs&quot;</span>]
    <span class="hljs-comment"># Now we train the model</span>
    <span class="hljs-keyword">for</span> epoch <span class="hljs-keyword">in</span> <span class="hljs-built_in">range</span>(num_epochs):
        model.train()
        <span class="hljs-keyword">for</span> step, batch <span class="hljs-keyword">in</span> <span class="hljs-built_in">enumerate</span>(train_dataloader):
            outputs = model(**batch)
            loss = outputs.loss
            accelerator.backward(loss)

            optimizer.step()
            optimizer.zero_grad()`,lang:"python",wrap:!1});var m=e(M,2);l(m,{code:"ZnJvbSUyMGFjY2VsZXJhdGUlMjBpbXBvcnQlMjBub3RlYm9va19sYXVuY2hlciUwQSUwQW5vdGVib29rX2xhdW5jaGVyKHRyYWluaW5nX2Z1bmN0aW9uKQ==",highlighted:`<span class="hljs-keyword">from</span> accelerate <span class="hljs-keyword">import</span> notebook_launcher

notebook_launcher(training_function)`,lang:"python",wrap:!1});var u=e(m,2);r(u,{children:(a,t)=>{var o=x();n(a,o)},$$slots:{default:!0}});var U=e(u,4);l(U,{code:"UHJvY2Vzc0V4aXRlZEV4Y2VwdGlvbiUzQSUyMHByb2Nlc3MlMjAwJTIwdGVybWluYXRlZCUyMHdpdGglMjBzaWduYWwlMjBTSUdTRUdW",highlighted:'<span class="hljs-attribute">ProcessExitedException</span>: process <span class="hljs-number">0</span> terminated <span class="hljs-keyword">with</span> <span class="hljs-keyword">signal</span><span class="hljs-string"> SIGSEGV</span>',lang:"",wrap:!1});var b=e(U,4);l(b,{code:"JTIzJTIwSW4lMjBhbm90aGVyJTIwSnVweXRlciUyMGNlbGwlMEFtb2RlbCUyMCUzRCUyMEF1dG9Nb2RlbEZvclNlcXVlbmNlQ2xhc3NpZmljYXRpb24uZnJvbV9wcmV0cmFpbmVkKCUyMmJlcnQtYmFzZS1jYXNlZCUyMiUyQyUyMG51bV9sYWJlbHMlM0QyKQ==",highlighted:`<span class="hljs-comment"># In another Jupyter cell</span>
model = AutoModelForSequenceClassification.from_pretrained(<span class="hljs-string">&quot;bert-base-cased&quot;</span>, num_labels=<span class="hljs-number">2</span>)`,lang:"python",wrap:!1});var w=e(b,2);l(w,{code:"JTJCJTIwZGVmJTIwdHJhaW5pbmdfZnVuY3Rpb24obW9kZWwpJTNBJTBBJTIwJTIwJTIwJTIwJTIwJTIwJTIzJTIwSW5pdGlhbGl6ZSUyMGFjY2VsZXJhdG9yJTBBJTIwJTIwJTIwJTIwJTIwJTIwYWNjZWxlcmF0b3IlMjAlM0QlMjBBY2NlbGVyYXRvcigpJTBBLSUyMCUyMCUyMCUyMCUyMG1vZGVsJTIwJTNEJTIwQXV0b01vZGVsRm9yU2VxdWVuY2VDbGFzc2lmaWNhdGlvbi5mcm9tX3ByZXRyYWluZWQoJTIyYmVydC1iYXNlLWNhc2VkJTIyJTJDJTIwbnVtX2xhYmVscyUzRDIpJTBBJTIwJTIwJTIwJTIwJTIwJTIwdHJhaW5fZGF0YWxvYWRlciUyQyUyMGV2YWxfZGF0YWxvYWRlciUyMCUzRCUyMGNyZWF0ZV9kYXRhbG9hZGVycyglMEElMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjAlMjB0cmFpbl9iYXRjaF9zaXplJTNEaHlwZXJwYXJhbWV0ZXJzJTVCJTIydHJhaW5fYmF0Y2hfc2l6ZSUyMiU1RCUyQyUyMGV2YWxfYmF0Y2hfc2l6ZSUzRGh5cGVycGFyYW1ldGVycyU1QiUyMmV2YWxfYmF0Y2hfc2l6ZSUyMiU1RCUwQSUyMCUyMCUyMCUyMCUyMCUyMCklMEElMjAlMjAuLi4=",highlighted:`<span class="hljs-addition">+ def training_function(model):</span>
      # Initialize accelerator
      accelerator = Accelerator()
<span class="hljs-deletion">-     model = AutoModelForSequenceClassification.from_pretrained(&quot;bert-base-cased&quot;, num_labels=2)</span>
      train_dataloader, eval_dataloader = create_dataloaders(
          train_batch_size=hyperparameters[&quot;train_batch_size&quot;], eval_batch_size=hyperparameters[&quot;eval_batch_size&quot;]
      )
  ...`,lang:"diff",wrap:!1});var f=e(w,4);l(f,{code:"JTIwJTIwZnJvbSUyMGFjY2VsZXJhdGUlMjBpbXBvcnQlMjBub3RlYm9va19sYXVuY2hlciUwQS0lMjBub3RlYm9va19sYXVuY2hlcih0cmFpbmluZ19mdW5jdGlvbiklMEElMkIlMjBub3RlYm9va19sYXVuY2hlcih0cmFpbmluZ19mdW5jdGlvbiUyQyUyMChtb2RlbCUyQykp",highlighted:`  from accelerate import notebook_launcher
<span class="hljs-deletion">- notebook_launcher(training_function)</span>
<span class="hljs-addition">+ notebook_launcher(training_function, (model,))</span>`,lang:"diff",wrap:!1});var g=e(f,2);r(g,{children:(a,t)=>{var o=N();n(a,o)},$$slots:{default:!0}});var j=e(g,2);i(j,{title:"Mixed Precision and Global Variables",local:"mixed-precision-and-global-variables",headingTag:"h2"});var J=e(j,8);l(J,{code:"YWNjZWxlcmF0b3IlMjAlM0QlMjBBY2NlbGVyYXRvcihtaXhlZF9wcmVjaXNpb24lM0QlMjJiZjE2JTIyKQ==",highlighted:'accelerator = Accelerator(mixed_precision=<span class="hljs-string">&quot;bf16&quot;</span>)',lang:"python",wrap:!1});var T=e(J,8);l(T,{code:"YWNjZWxlcmF0b3IlMjAlM0QlMjBBY2NlbGVyYXRvcihtaXhlZF9wcmVjaXNpb24lM0QlMjJiZjE2JTIyJTJDJTIwZG93bmNhc3RfYmYxNiUzRFRydWUp",highlighted:'accelerator = Accelerator(mixed_precision=<span class="hljs-string">&quot;bf16&quot;</span>, downcast_bf16=<span class="hljs-literal">True</span>)',lang:"python",wrap:!1});var C=e(T,4);i(C,{title:"Training Times on TPUs",local:"training-times-on-tpus",headingTag:"h2"});var G=e(C,6);r(G,{children:(a,t)=>{var o=E();n(a,o)},$$slots:{default:!0}});var v=e(G,2);W(v,{source:"https://github.com/huggingface/accelerate/blob/main/docs/source/concept_guides/training_tpu.md"}),F(2),n(Y,c),B()}export{D as component};