استخدام SGD المحلي مع Accelerate

تُعدّ SGD المحلي تقنيةً للتدريب الموزّع لا يتم فيها مزامنة التدرّجات في كل خطوة. وبالتالي، يقوم كل عملية بتحديث نسخته الخاصة من أوزان النموذج، وبعد عدد معيّن من الخطوات يتم مزامنة هذه الأوزان عن طريق أخذ المتوسط عبر جميع العمليات. يحسّن هذا الكفاءة في التواصل ويمكن أن يؤدي إلى تسريع كبير في التدريب، خاصةً عندما تفتقر الحاسوب إلى وصلة أسرع مثل NVLink. على عكس تراكم التدرّجات (حيث يتطلب تحسين كفاءة التواصل زيادة حجم الدفعة الفعّال)، لا يتطلب SGD المحلي تغيير حجم الدفعة أو معدل التعلم / الجدول الزمني. ومع ذلك، إذا لزم الأمر، يمكن دمج SGD المحلي مع تراكم التدرّجات أيضًا.

في هذا الشرح التعليمي، ستتعلم كيفية إعداد Local SGD Accelerate بسرعة. مقارنةً بإعداد Accelerate القياسي، يتطلب هذا الأمر سطرين إضافيين فقط من الكود.

سيستخدم هذا المثال حلقة تدريب PyTorch بسيطة جداً تقوم بتراكم التدرجات كل دفعتين:

device = "cuda"
model.to(device)

gradient_accumulation_steps = 2

for index, batch in enumerate(training_dataloader):
    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    loss = loss / gradient_accumulation_steps
    loss.backward()
    if (index + 1) % gradient_accumulation_steps == 0:
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()

تحويله إلى Accelerate

أولاً، سيتم تحويل الكود المعروض سابقاً لاستخدام Accelerate دون استخدام LocalSGD أو مساعد تراكم التدرجات:

+ from accelerate import Accelerator
+ accelerator = Accelerator()

+ model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+     model, optimizer, training_dataloader, scheduler
+ )

  for index, batch in enumerate(training_dataloader):
      inputs, targets = batch
-     inputs = inputs.to(device)
-     targets = targets.to(device)
      outputs = model(inputs)
      loss = loss_function(outputs, targets)
      loss = loss / gradient_accumulation_steps
+     accelerator.backward(loss)
      if (index+1) % gradient_accumulation_steps == 0:
          optimizer.step()
          scheduler.step()

السماح لـ Accelerate بالتعامل مع مزامنة النموذج

كل ما تبقّى الآن هو أن نترك لـ Accelerate مهمة مزامنة معاملات النموذج و تراكم التدرجات نيابةً عنا. وبسبب البساطة، لنفترض أننا بحاجة إلى المزامنة كل 8 خطوات. يتم تحقيق ذلك بإضافة عبارة with LocalSGD واحدة واستدعاء local_sgd.step() واحد بعد كل خطوة للمحسّن:

+local_sgd_steps=8

+with LocalSGD(accelerator=accelerator, model=model, local_sgd_steps=8, enabled=True) as local_sgd:
    for batch in training_dataloader:
        with accelerator.accumulate(model):
            inputs, targets = batch
            outputs = model(inputs)
            loss = loss_function(outputs, targets)
            accelerator.backward(loss)
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()
+           local_sgd.step()

في الخلفية، يُعطّل كود Local SGD مزامنة التدرجات تلقائياً (لكن التراكم لا يزال يعمل كما هو متوقع!). بدلاً من ذلك، يقوم بمتوسط معاملات النموذج كل local_sgd_steps خطوة (وكذلك في نهاية حلقة التدريب).

القيود

يعمل التنفيذ الحالي فقط مع التدريب الأساسي متعدد وحدات معالجة الرسومات (أو متعدد وحدات المعالجة المركزية) دون استخدام، على سبيل المثال، DeepSpeed.

المراجع

على الرغم من أننا لا نعلم بالتحديد الأصول الحقيقية لهذا النهج البسيط، إلا أن فكرة SGD المحلية قديمة نسبياً وتعود على الأقل إلى:

Zhang, J., De Sa, C., Mitliagkas, I., & Ré, C. (2016). Parallel SGD: When does averaging help?. arXiv preprint arXiv:1606.07365.

نُعزّي مصطلح Local SGD إلى الورقة البحثية التالية (لكن قد تكون هناك مراجع سابقة لا نعلم بها).

Stich, Sebastian Urban. “Local SGD Converges Fast and Communicates Little.” ICLR 2019-International Conference on Learning Representations. No. CONF. 2019.