آموزش مدل های غول پیکر با مقیاس گذاری نزدیک به خطی با استفاده از موازی سازی داده های خرد شده در Amazon SageMaker

در جستجوی دقت برتر، مدل‌های یادگیری عمیق در زمینه‌هایی مانند پردازش زبان طبیعی و بینایی رایانه‌ای در چند سال گذشته به‌طور قابل توجهی رشد کرده‌اند و اغلب در ده‌ها تا صدها میلیارد پارامتر محاسبه می‌شوند. آموزش این مدل های غول پیکر چالش برانگیز است و به استراتژی های توزیع پیچیده ای نیاز دارد. دانشمندان داده و مهندسان یادگیری ماشین دائماً به دنبال بهترین راه برای بهینه‌سازی محاسبات آموزشی خود هستند، اما با هزینه‌های ارتباطی که می‌تواند همراه با اندازه کلی خوشه افزایش یابد، دست و پنجه نرم می‌کنند.

به همین دلیل است که ما اخیرا راه اندازی کردیم sموازی سازی داده های سخت on آمازون SageMaker، یک تکنیک آموزشی توزیع شده ذخیره سازی حافظه جدید در کتابخانه موازی مدل SageMaker (SMP).. موازی سازی داده های خرد شده به طور هدفمند برای مدل های در مقیاس شدید ساخته شده است و از آمازون در داخل استفاده می کند MiCS فناوری زیر کاپوت، یک تلاش علمی برای به حداقل رساندن مقیاس ارتباطی با پایین آوردن هزینه های ارتباطی گران قیمت که ریشه در جمع آوری پارامترها و همگام سازی گرادیان دارد. با پارامتر 30B مدل GPT-2 با طول توالی 2048، این ویژگی جدید به 141 TFLOP دست یافت که در مقایسه با DeepSpeed ​​ZeRO-39.7 3 درصد سرعت داشت. برای یک مدل 10B GPT-2 با طول دنباله 512، این ویژگی جدید همچنین به 564 نمونه در ثانیه دست یافت که 13.9٪ سرعت در مقایسه با PyTorch's Fully Sharded Data Parallel (FSDP) دارد. به یاد داشته باشید که در آموزش مدل های غول پیکر، هر درصد افزایش سرعت به دلار پس انداز و بهره وری در تیم شما تبدیل می شود.

در این پست وبلاگ، ابتدا نگاهی دقیق تر به تمایزهای کلیدی موازی سازی داده های خرد شده و زمان استفاده از آن خواهیم داشت. سپس، نحوه آموزش مدل GPT-30 با پارامتر 2B را در SageMaker به راحتی با این ویژگی جدید خواهید آموخت. در نهایت، ما عملکرد را با سایر گزینه‌های منبع باز مقایسه می‌کنیم، به ویژه عملکرد بهتر از DeepSpeed ​​ZeRO تا 39.7٪ در 256 GPU.

نحوه عملکرد موازی داده های خرد شده و زمان استفاده از آن

قبل از معرفی موازی سازی داده های خرد شده، اجازه دهید به خانواده تکنیک های گسترده تر آن نگاه کنیم. رویکردهای آموزشی توزیع شده اخیر برای مدل‌های بزرگ به پارادایم منتقل شده‌اند که در آن پارامترهای مدل، گرادیان‌ها و حالت‌های بهینه‌ساز در میان گره‌های موازی داده به اشتراک گذاشته می‌شوند. برخلاف Pipeline Parallelism که پیچیدگی ذاتی انتخاب لایه ها برای پارتیشن بندی در بین دستگاه ها را دارد، به خصوص زمانی که چارچوب شما پشتیبانی نمی کند. تقسیم خودکار مدل، این پارادایم به زیبایی سادگی موازی سازی داده ها را حفظ می کند، در حالی که محدودیت موازی سازی داده را که در آن یک مدل باید در یک GPU واحد قرار گیرد حذف می کند.

در چارچوب‌های موجود که تحت این پارادایم قرار می‌گیرند، به‌ویژه DeepSpeed ​​ZeRO-3 و PyTorch's FSDP که از FairScale بالادستی شده است، حالت‌های مدل در همه جا تقسیم می‌شوند. تمام پردازنده‌های گرافیکی، استراتژی‌ای که مصرف حافظه را در هر GPU به قیمت تحمیل هزینه‌های ارتباطی زیاد کاهش می‌دهد که با اندازه خوشه افزایش می‌یابد و در نتیجه باعث می‌شود مقیاس‌پذیری به میزان قابل توجهی در مقیاس کاهش یابد. در مقابل، موازی سازی داده های خرد شده در مدل پارتیشن های کتابخانه SMP در الف بیان می شود مقیاس آگاه با پارتیشن بندی هر کپی از حالت های مدل فقط در داخل یک زیر مجموعه از پردازنده های گرافیکی

بیایید از نزدیک نگاه کنیم پارتیشن بندی مدل آگاه از مقیاس در MiCS، فناوری اصلی پشت موازی داده های خرد شده. شهود پشت این طراحی این است که پارتیشن بندی حالت های آموزشی در کل گروه موازی داده ممکن است برای آموزش مدلی با ده ها میلیارد پارامتر مورد نیاز نباشد. به عنوان مثال، 8 پردازنده گرافیکی V100 (هر کدام 32 گیگابایت) برای نگهداری حالت های مدل یک مدل با پارامتر 10B کافی است که هنگام آموزش با بهینه ساز Adam با استفاده از دقت ترکیبی، به حدود 200 گیگابایت حافظه نیاز دارد. با محدود کردن یک کپی کامل از حالت های مدل در کوچکترین زیرمجموعه‌ای از پردازنده‌های گرافیکی، می‌توانیم به طور موثر مقیاس سربار ارتباط را در مقایسه با DeepSpeed ​​و PyTorch FSDP کاهش دهیم. موازی داده‌های خرد شده همچنین از تکنیک‌های دیگر در MiCS مانند ارتباطات سلسله مراتبی و همگام‌سازی گرادیان 2-hop استفاده می‌کند. برای اطلاعات بیشتر، بررسی کنید مقیاس بندی تقریباً خطی آموزش مدل غول پیکر در AWS or MiCS: مقیاس گذاری نزدیک به خطی برای آموزش مدل غول پیکر در ابر عمومی.

اکنون، چگونه می‌دانید که چه زمانی داده‌های خرد شده را به موازات سایر تکنیک‌های آموزشی توزیع‌شده انتخاب کنید؟ قانون کلی این است که اگر مدل شما کمتر از 1 میلیارد پارامتر داشته باشد و بتواند در حافظه GPU جای بگیرد، کتابخانه موازی داده SageMaker or کامپایلر آموزش SageMaker می تواند برای شما کافی باشد اگر زبان یا مدل‌های بینایی کامپیوتری بزرگ‌تری دارید، پیشنهاد ما این است که آن را با تکنیک موازی‌سازی داده‌های خرد شده همراه با ایست بازرسی فعال سازی و بارگذاری فعال سازی در کتابخانه موازی مدل SageMaker ابتدا قبل از تکنیک های دیگری مانند موازی تانسور یا موازی خط لوله.

استفاده از موازی سازی داده های خرد شده برای آموزش GPT-2 در Amazon SageMaker

بیایید اکنون بیاموزیم که چگونه یک مدل GPT-2 را با داده های خرد شده موازی آموزش دهیم، با SMP که پیچیدگی را برای شما محصور می کند. این دفترچه آموزشی کامل شما را در کل فرآیند از پردازش داده ها، تعریف و ارسال مشاغل آموزشی گرفته تا نظارت بر گزارش های آموزشی راهنمایی می کند. آنچه در زیر می آید مروری کوتاه است که مراحل کلیدی استفاده از این ویژگی را برجسته می کند.

1. شروع کنید

موازی سازی داده های خرد شده در PyTorch نسخه 1.12.0+ موجود است و با هر دو FP16 و BF16 کار می کند. ساده ترین راه برای استفاده از کتابخانه SMP از طریق یک کانتینر آموزش عمیق AWS از پیش ساخته شده برای PyTorch است. با این حال، اگر می خواهید کانتینر داکر خود را بیاورید، می توانید به آن مراجعه کنید با کتابخانه موازی مدل توزیع شده SageMaker کانتینر Docker خود را ایجاد کنید. برای شروع، فالو کنید یک اسکریپت آموزشی PyTorch را تغییر دهید برای تطبیق APIهای SMP در اسکریپت آموزشی شما. در این بخش، ما فقط چند مرحله اصلی را با قطعه کد از اسکریپت آموزشی آماده برای استفاده فراخوانی می کنیم. train_gpt_simple.py. می توانید نظرات را در اسکریپت دنبال کنید و سند API برای کسب اطلاعات بیشتر در مورد محل استفاده از SMP API.

ابتدا کتابخانه را با فراخوانی وارد کرده و مقداردهی اولیه کنید smdistributed.modelparallel.torch.init() در ابتدای متن آموزشی:

import smdistributed.modelparallel.torch as smp

smp.init(smp_config)

دوم، مدلی که قرار است با آن پارتیشن بندی شود را بپیچید smdistributed.modelparallel.torch.DistributedModel و از برگردانده استفاده کنید DistributedModel شی به جلو:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_config(model_config)
model = smp.DistributedModel(model, trace_device="gpu", backward_passes_per_step=args.gradient_accumulation)

بهینه ساز را با smdistributed.modelparallel.torch.DistributedOptimizer برای ذخیره و بارگیری حالات بهینه ساز.

from torch import optim

optimizer = optim.Adam(
    param_groups, betas=(args.beta1, args.beta2), lr=args.lr, weight_decay=args.weight_decay
)

optimizer = smp.DistributedOptimizer(
        optimizer, 
        static_loss_scale=None, 
        dynamic_loss_scale=True,
        dynamic_loss_args={"scale_window": 1000, "min_scale": 1, "delayed_shift": 2},
        )

منطق جلو و عقب را در یک تابع step قرار دهید و آن را با آن تزئین کنید smdistributed.modelparallel.torch.step.  هر محاسباتی که در داخل تعریف شده است smp.step-decorated تابع به صورت توزیع شده اجرا می شود.

@smp.step
def train_step(model, optimizer, input_ids, attention_mask, args):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    model.backward(loss)

    return loss

@smp.step
def test_step(model, input_ids, attention_mask):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    
    return loss

2. مجموعه داده را آماده کنید

ما با استفاده از متن باز وب مجموعه داده ای است که در این مثال استفاده می کنیم. نوت بوک از اسکریپت استفاده می کند data_prep_512.py برای دانلود و پیش پردازش مجموعه داده. شما همچنین می توانید با سایر مجموعه داده ها با تغییر آموزش دهید data_pipeline.py. هنگامی که با مجموعه داده ها و مدل های بزرگ سروکار دارید، می توانید با استفاده از داده های ذخیره شده در کار آموزش را سرعت بخشید آمازون FSx برای Luster، که یک سیستم فایل با کارایی بالا را به صورت بومی با آن ادغام می کند سرویس ذخیره سازی ساده آمازون (S3). لطفا دستورالعمل ها را از کانال ورودی داده را برای استفاده از Amazon FSx برای Luster پیکربندی کنید برای راهنمایی در مورد تنظیم یک سیستم فایل FSx Luster به عنوان کانال ورودی داده.

3. مشاغل آموزشی را شروع کنید

این مرحله فرض می کند که قبلاً این کار را کرده اید اسکریپت آموزشی شما را اصلاح کرد و مجموعه داده را همانطور که در قسمت های قبل ذکر شد آماده کرد. به موازی سازی داده های خرد شده را فعال کنید، به سادگی آن را تنظیم کنید sharded_data_parallel_degree در برآوردگر PyTorch. در این آموزش تنظیم کردیم sharded_data_parallel_degree=128 و instace_count=32 برای گره‌های p4d.24xlarge، که نشان می‌دهد که حالت‌های مدل در 128 GPU از مجموع 256 GPU تقسیم می‌شوند. بر اساس این مقدار انتخاب شده، سپس SMP به طور خودکار درجه موازی داده ها را روی 2 تنظیم می کند (زیرا 256/128=2)، به این معنی که ما دو کپی برای موازی سازی داده خواهیم داشت. یک قانون کلی برای انتخاب یک مقدار ایده آل برای sharded_data_parallel_degree به ازای هر 3B پارامتر مدل، یک گره دیگر به گروه اشتراک‌گذاری اضافه می‌شود. در این آموزش، اندازه مدل ما 30B است، بنابراین باید حداقل از 10 گره برای اشتراک گذاری استفاده کنیم. و از آنجایی که 16 گره (128 GPU) کوچکترین توان 2 بالاتر از آستانه است، ما تنظیم کردیم sharded_data_parallel_degree=128.

برای ایست بازرسی، مجموعه ای از ابزارهای بازرسی را نیز در اختیار شما قرار می دهیم sharded_data_parallel_checkpoint.py ، از جمله ابزاری برای بازسازی کامل state_dict برای موارد استفاده پیشرفته در نهایت، می‌توانیم یک کار آموزشی توزیع‌شده را با فراخوانی fit() در برآوردگر راه‌اندازی کنیم.

smp_estimator = PyTorch(
    entry_point="train_gpt_simple.py",
    instance_type="ml.p4d.24xlarge",
    source_dir=os.getcwd(),
    volume_size=500,
    instance_count=32,
    distribution={
        "mpi": {
            "enabled": True,
            "processes_per_host": processes_per_host,
            "custom_mpi_options": mpioptions,
        },
        "smdistributed": {
            "modelparallel": {
                "enabled": True,
                "parameters": {
                    "ddp": True,
                    "skip_tracing": True,
                    "delayed_parameter_initialization": True,
                    "offload_activations": True,
                    "activation_loading_horizon": 4,
                    # To enable sharded data parallelism.
                    # Here we shard model states across 128 GPUs. 
                    "sharded_data_parallel_degree": 128, 
                    "fp16": False,
                    "bf16": True,
                    # This is to disable pipeline parallelism.
                    "partitions": 1,
                },
            }
        },
    },
    framework_version="1.12",
    py_version="py38",
    hyperparameters=hyperparameters,
    checkpoint_s3_uri=checkpoint_s3_uri if not use_fsx else None,
    checkpoint_local_path=hyperparameters["checkpoint-dir"] if use_fsx else None,
    ...
)

smp_estimator.fit(inputs=data_channels)

4. نظارت بر مشاغل آموزشی

می‌توانید به گزارش‌های آموزشی دسترسی داشته باشید و میزان استفاده از GPU و حافظه را ردیابی کنید CloudWatch آمازون. حتماً به گزارش های "algo-1" نگاه کنید زیرا این گره اصلی است که جریان خروجی آن دارای گزارش های شغلی آموزشی از همه نمونه ها است.

ارزیابی عملکرد

ما موازی سازی داده های خرد شده را در کتابخانه SMP در هر دو 16 و 32 گره p4d.24xlarge به ترتیب برای طول دنباله 512 و 2048 محک زدیم. مدل GPT30 با پارامتر 2B برای استفاده از عرض پنهان 7168، 48 لایه و 64 هد پیکربندی شده است. شما می توانید دقیقاً همان پیکربندی را که در آن طول دنباله 2048 است با تنظیم استفاده کنید model_config = "gpt2-30b" در دفترچه آموزشی با این تنظیم، SMP به 73.52 نمونه در ثانیه دست یافت که در مقایسه با DeepSpeed ​​ZeRO-39.7 3 درصد سرعت داشت. اگر اندازه توکن شما 500 میلیارد است، این افزایش سرعت به معنای نزدیک به 367 ساعت صرفه جویی در گره های p4d.24xlarge است که معادل بیش از 12,000 دلار بودجه صرفه جویی در هر آموزش است! جدول زیر نتایج محک ما را خلاصه می کند.

1/ برای پیکربندی هر مدل، ویژگی‌ها، مراحل و پیکربندی‌های مختلفی را در DeepSpeed ​​ZeRO آزمایش کردیم و یکی را انتخاب کردیم که بهترین توان عملیاتی را به عنوان خط پایه DeepSpeed ​​ارائه می‌دهد. معیار اجرا شد ابر محاسبه الاستیک آمازون (آمازون EC2). 2/ این نتایج به مجموعه‌های ارتباطی بهبودیافته بهینه‌سازی شده برای AWS متکی است که به زودی در دسترس قرار خواهند گرفت. 3/ زمان آموزش از سرعت بر اساس تعداد توکن های پردازش شده پیش بینی می شود.

به طور خلاصه، ما در مقایسه با DeepSpeed ​​در طیف وسیعی از مدل‌ها و پیکربندی‌ها، توان عملیاتی بالاتری را با موازی‌سازی داده‌های خرد شده در SMP مشاهده کردیم. این ویژگی جدید همچنین کارایی حافظه بهتری را در مقایسه با DeepSpeed ​​نشان می‌دهد، و SMP را قادر می‌سازد تا با اندازه دسته‌ای بزرگ‌تر مطابقت داشته باشد و سطح انباشتگی گرادیان مورد نیاز برای تناسب با اندازه دسته‌ای خاص را کاهش دهد.

نتیجه

در این پست، یک تکنیک آموزشی توزیع‌شده جدید - موازی‌سازی داده‌های خرد شده - و چگونگی سرعت بخشیدن به آموزش مدل‌های غول‌پیکر با مقیاس‌گذاری تقریباً خطی در Amazon SageMaker را معرفی کردیم. ما همچنین نحوه آموزش یک مدل GPT-2 را با تکنیک جدید به دنبال این موضوع توضیح دادیم مثال کامل. می توانید دنبال کنید آمازون SageMaker مخزن GitHub را مثال می‌زند برای ردیابی تمام نمونه های موازی مدل SageMaker یا شرکت در آینده ما کارگاه های آموزشی توزیع کرد. برای کسب اطلاعات بیشتر در مورد موازی سازی داده های خرد شده، لطفاً به اینجا مراجعه کنید مستندات.

درباره نویسندگان

امیلی وبر درست پس از راه اندازی SageMaker به AWS ملحق شد و از آن زمان تلاش کرده است تا در مورد آن به جهان بگوید! غیر از ایجاد تجربیات جدید ML برای مشتریان، امیلی از مراقبه و مطالعه بودیسم تبتی لذت می برد.

می توان کاراکوس یک دانشمند ارشد کاربردی در AWS است که یادگیری عمیق توزیع شده در مقیاس بزرگ را در AWS بهینه می کند. علایق تحقیقاتی او شامل یادگیری عمیق، بهینه سازی توزیع شده، سیستم های توزیع شده و نظریه اطلاعات است. خارج از محل کار، او از دوچرخه سواری، مسافرت، مطالعه و یادگیری لذت می برد.

راهول هویل گل مهندس نرم افزار ارشد در AWS است. او روی سیستم‌های یادگیری عمیق توزیع شده کار می‌کند تا آموزش مدل‌های یادگیری عمیق بزرگ را در فضای ابری آسان و کارآمد کند. او در اوقات فراغت خود از عکاسی، دوچرخه سواری و باغبانی لذت می برد.

سوهیت کودگوله یک مهندس توسعه نرم افزار با گروه هوش مصنوعی AWS است که بر روی چارچوب های یادگیری عمیق کار می کند. در اوقات فراغت از پیاده روی، مسافرت و آشپزی لذت می برد.

ارین هو یک مدیر محصول برای یادگیری عمیق AWS است. او روی محصولاتی کار می کند که آموزش مدل های یادگیری عمیق را در AWS برای مشتریان آسان تر می کند. برای تفریح ​​در خارج از خانه، او از پیاده روی و اسکی لذت می برد.