تحليل أداء TorchNeuron (الجزء 1): تحليل الأداء على مستوى النظام على AWS Trainium.

يمتلك كثير من ممارسي تعلم الآلة حدسًا معينًا حول كيفية جعل النماذج تعمل بسرعة، لكن قلة منهم يعرفون كيف يحللون سبب بطء النموذج. في الجزء 1 من هذه السلسلة، ستتعلم كيف تحلل أداء الاستدلال (inference) بشكل منهجي وتحدد أين يُستهلك الوقت.

يبدأ هذا المنشور من الصفر ويفترض عدم وجود خبرة سابقة في تحليل الأداء. تستخدم هذه الجولة العملية AWS Trainium، وهي شريحة ذكاء اصطناعي مخصصة لأحمال تدريب واستدلال التعلم العميق، و TorchNeuron، وهو الواجهة الخلفية الأصلية مفتوحة المصدر لـ PyTorch. يدعم TorchNeuron سير عمل PyTorch القياسي على Trainium، بما في ذلك التنفيذ الفوري (eager execution)، torch.compileوواجهات PyTorch القياسية، مع الحد الأدنى من تغييرات الشيفرة. بالإضافة إلى ذلك، يتكامل TorchNeuron مع Neuron Kernel Interface (NKI) لتمكين العمليات المخصصة وتحسين الأداء لعتاد Trainium. بنهاية المقال، ينبغي أن تعرف:

• كيفية تهيئة torch.profiler مع الواجهة الخلفية Neuron والتقاط تتبع نظيف.
• كيفية قراءة جدول key_averages() لتحديد العمليات التي تستهلك وقت جهاز Neuron ومقارنة زمن Self CPU الخاص بها.
• كيفية تتبع عملية واحدة عبر Neuron Explorer، من إرسالها في الإطار (framework dispatch) مرورًا بجدولتها في بيئة التشغيل (runtime scheduling) وصولًا إلى تنفيذها على العتاد في NeuronCore.
• كيف تتغير هذه الإشارات عبر ثلاثة تطبيقات للعملية نفسها.

Neuron كواجهة خلفية لـ PyTorch

قبل أن تحلل أداء حمل عمل، من المفيد أن تفهم كيف يتكامل Neuron مع PyTorch.

تستخدم الواجهة الخلفية Neuron نقطة امتداد PrivateUse1 الخاصة بـ PyTorch. تتيح هذه الآلية لمسرّعٍ ما التكامل مع الإطار دون إضافة مفتاح إرسال (dispatch key) مخصص مثل cuda. تعمل حزمة torch_neuronx مع PyTorch الأصلي (upstream) لتسجيل Neuron كواجهة خلفية أصلية لشيفرة PyTorch. للحصول على تفاصيل حول آلية الامتداد هذه، راجع Facilitating New Backend Integration by PrivateUse1.

إعداد أداة تحليل الأداء (Profiler)

تستخدم كائنين لتحليل أداء حمل عمل Neuron:

• NeuronConfig يتحكم في بيانات Neuron ووحدات NeuronCore التي سيُلتقط تدفقها ومكان كتابة المخرجات. تمرّره إلى torch.profiler.profile بصيغة experimental_config.
• NeuronProfiler تأخذ NeuronConfig وتوفر استدعاءً رجعيًا (callback) من نوع export_trace يكتب الاستدعاء الرجعي تتبع الإطار (framework trace)، trace.json، إلى جانب تتبعات نظام Neuron.

nc = NeuronConfig(
    modes=[ProfileMode.DEVICE, ProfileMode.RUNTIME],
    profile_output_dir=OUTPUT_DIR,
    capture_enabled_for_nc='0',
    correlate_device_time=True,
)
exporter = NeuronProfiler(nc)

الـ modes تحدد تدفقات البيانات المطلوب التقاطها، مثل تعليمات الجهاز وآثار التشغيل. أما الـ capture_enabled_for_nc فتحدد NeuronCores المطلوب تتبعها. يؤدي تعيين correlate_device_time=True إلى نشر زمن الجهاز من التتبع إلى جدول الـ key_averages() حيث يمكنك مقارنة زمن CPU وزمن جهاز Neuron. للحصول على المرجع الكامل للمعاملات وقيم الـ ProfileMode المتاحة، راجع وثائق NeuronConfig أوضاع التنميط (Profile Modes).

الأدوات: Neuron Explorer

Neuron Explorer هو عارض قائم على المتصفح لتتبع النظام والخط الزمني للجهاز وتفصيل كل محرك. يمكنك الوصول إليه عبر نفق SSH أو إضافة Visual Studio Code. في هذا المنشور، تستخدم Explorer لتأكيد ما يبلغ عنه key_averages() رقميًا وفحص ما يحدث بين عمليات الإرسال.

بعد التنميط، يحتوي مجلد المخرجات على تتبع النظام، وعند التقاط بيانات الجهاز، ملفات الـ .neff والـ .ntff المتطابقة. يحتوي تنسيق ملف Neuron القابل للتنفيذ (NEFF) على الملف التنفيذي لـ Neuron المترجم، بينما يحتوي ملف NTFF المطابق على تتبع تعليمات العتاد الخاص به. تشترك الملفات في تجزئة رقمية تميزهما كزوج. لفتح التنميط في Neuron Explorer، ارفع مجلد المخرجات الكامل لتنميط النظام أو التنميط المدمج. أما بالنسبة لتنميط الجهاز فقط، فارفع ملفي الـ .neff والـ .ntff المتطابقين فرديًا. راجع الخطوة 5: ارفع ملف التنميط الخاص بك للحصول على تعليمات واجهة المستخدم وسطر الأوامر.

تنميط RMSNorm

في هذا التحليل الأول، تنمّط RMSNorm، وهي عملية صغيرة ذات تسلسل واضح من عمليات الحساب والتقليص.

تنمّط ثلاث تطبيقات لنفس العملية:

1. RMSNorm بسيط بعمليات PyTorch فردية: الضرب، والمتوسط، والإضافة، والجذر التربيعي المقلوب، والضرب.
2. torch.rms_norm()، وهي الدالة المدمجة في PyTorch ومخفضة إلى نواة Neuron منصهرة واحدة.
3. RMSNorm بسيط باستخدام torch.compileبنفس العمليات المستخدمة في التطبيق الأول.

تستخدم النصوص البرمجية نفس المدخلات ونفس إعدادات Profiler ونفس الجدولة بحيث تكون النتائج قابلة للمقارنة مباشرة.

جميع النصوص البرمجية اللازمة لإعادة إنتاج هذا العمل متوفرة في torch-neuronx مجلد الأمثلة examples directory.

الملف الشخصي 1: عمليات PyTorch البسيطة

ابدأ بكتابة RMSNorm باستخدام عمليات PyTorch الفردية. كل سطر يقابل عملية واحدة أو أكثر من عمليات ATen ويمكن أن ينشئ موترات وسيطة.

device = torch.device('neuron')

x = torch.randn(4, 512, dtype=torch.bfloat16, device=device)
weight = torch.ones(512, dtype=torch.bfloat16, device=device)

def rms_norm_fn(t):
    x_sq = t * t                                    # square
    mean_sq = x_sq.mean(dim=-1, keepdim=True)       # mean of squares
    rms = torch.rsqrt(mean_sq + 1e-6)               # 1 / sqrt(mean_sq + eps)
    return t * rms * weight                         # normalize + scale

استدعِ profiler بوسيطين مطلوبين لخادم Neuron الخلفي:

1. قم بتضمين ProfilerActivity.PrivateUse1 لجمع النشاط من جهاز Neuron.
2. مرر NeuronConfig بقيمة experimental_config لتمكين ميزات Kineto الخاصة بـ Neuron.

تستخدم الجدولة wait=2, warmup=1و active=1 عبر أربع استدعاءات لـ prof.step() يسجل هذا تكرار تسخين واحد دون تضمين تكاليف الإعداد الأولي.

استدعِ torch.neuron.synchronize() داخل حلقة profiler. بدون هذه المزامنة، قد يصدر profiler البيانات قبل انتهاء الجهاز ويترك عمود وقت Neuron فارغًا.

with torch.profiler.profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.PrivateUse1],
    experimental_config=nc,
    record_shapes=False,
    on_trace_ready=exporter.export_trace,
    schedule=torch.profiler.schedule(wait=2, warmup=1, active=1),
) as prof:
    for step in range(4):
        with torch.no_grad():
            _ = rms_norm_fn(x)
        torch.neuron.synchronize()
        prof.step()

بعد انتهاء عملية التصنيف (profiling)، افحص جدول key_averages() رتب حسب self_privateuse1_time_total لوضع العمليات ذات التكلفة الأعلى على الجهاز في المقدمة. الجدول المعروض هنا مُفلتر ليشمل عمليات RMSNorm التي تُبلغ عن وقت جهاز Neuron. الناتج الكامل لـ key_averages() يتضمن أيضًا أحداثًا على جانب CPU مثل التخصيصات وعمليات النقل من المضيف إلى الجهاز.

print(prof.key_averages().table(sort_by="self_privateuse1_time_total", row_limit=4))

يوضح جدول المتوسطات الرئيسية أن ثلاث استدعاءات لـ aten::mul تمثل الحصة الأكبر من وقت جهاز Neuron.

يُظهر الجدول عمليات ATen التي جرى تنفيذها، وعدد مرات تنفيذها، ومقدار الوقت الذي أمضته على CPU وجهاز Neuron. aten::mul له أعلى إجمالي وقت على الجهاز لأن الدالة تؤثر ثلاث عمليات ضرب.

الإجمالي الذاتي لـ NEURON: 2.355 ms

كل صف يتطابق مباشرة مع الدالة. aten::mul يعمل لمدة x * x, x * rms، وتضرب القيمة النهائية في weight. aten::mean يقلل البعد الأخير. aten::add يضيف قيمة epsilon ويطلق عملية نسخ من المضيف إلى الجهاز لأن العدد القياسي يبدأ على وحدة المعالجة المركزية. aten::rsqrt يحسب الجذر التربيعي المتبادل.

يتتبع Neuron Explorer عملية aten::mul من إرسال الإطار عبر جدولة وقت التشغيل وصولاً إلى التنفيذ على مسارات أجهزة Neuron.

اقرأ التتبع في Explorer

في Explorer، يتدفق الزمن من اليسار إلى اليمين وتتدفق دورة حياة الحدث من الأعلى إلى الأسفل. تتبع كل عملية ATen النمط نفسه: الإنشاء في مسار الإطار، والجدولة في مسار وقت التشغيل بصيغة nrta_execute_schedule، والتنفيذ على الجهاز في مسار العتاد بصيغة nc_exec_running.

تنظم وعرض تتبع النظام الأحداث في مسارات. في وضع التجميع الافتراضي (Default)، تتبع التسميات النمط instance/PID/subsystem/identifier:

في الوضع الفوري (eager mode)، تعمل كل عملية بشكل مستقل. استخدم عارض سلسلة التبعيات لتتبع aten::mul.Tensor من طبقة الإطار مروراً بجدولة وقت التشغيل وصولاً إلى التنفيذ على الجهاز.

يسلط عارض سلسلة التبعيات الضوء على دورة حياة عملية واحدة عبر مسارات الإطار ووقت التشغيل والعتاد.

التنميط 2: torch.rms_norm

استبدل العمليات الفردية بدالة rms_norm المدمجة في PyTorch. يقوم Neuron بتحويل العملية إلى نواة (kernel) واحدة مدمجة:

def rms_norm_fn(t):
    return torch.rms_norm(t, [512], weight)

يظهر RMSNorm المدمج كعملية جهاز واحدة في جدول المتوسطات الرئيسية.

إجمالي Self NEURON: 403.3 μs

يتبع RMSNorm المدمج مسارًا واحدًا من الإطار إلى العتاد بدلاً من إرسال كل عملية حسابية على حدة.

قارن مسار التخفيض (Lowering Path)

في التنفيذ البسيط، تعبر كل عملية حدود الإرسال بشكل مستقل وتتبع مسارها الخاص من الإطار إلى بيئة التشغيل إلى العتاد. في التنفيذ المدمج، تغطي استدعاء بايثون واحد العملية الحسابية كاملة. يحتوي التتبع على عملية إرسال واحدة وحدث عتادي واحد، دون فجوات بين العمليات المكونة.

الملف الشخصي 3: PyTorch بسيط مع torch.compile

في المثال الأخير، قم بترجمة دالة RMSNorm البسيطة الأصلية باستخدام torch.compile.

مع @torch.compile(backend="neuron")، يلتقط TorchDynamo شيفرة بايثون البايتية كرسم بياني FX وتقوم AOTAutograd بتخفيضه إلى عمليات ATen. يتلقى الواجهة الخلفية لـ Neuron الرسم البياني الكامل بدلاً من إرسال كل عملية على حدة. يدمج العمليات، ويحسّن الوصول إلى الذاكرة، ويترجم الرسم البياني إلى قطعة أثرية NEFF واحدة. تعمل القطعة الأثرية كعملية إرسال عتادية واحدة.

@torch.compile(backend='neuron', fullgraph=True)
def naive_compiled(t):
    x_sq = t * t
    mean_sq = x_sq.mean(dim=-1, keepdim=True)
    rms = torch.rsqrt(mean_sq + 1e-6)
    return t * rms * weight

تظهر الدالة المترجمة كمنطقة واحدة باسم Torch-Compiled Region في جدول المتوسطات الرئيسية.

إجمالي وقت NEURON الذاتي: 375.5 μs

يتبع الرسم البياني المترجم مسارًا واحدًا عبر جدولة بيئة التشغيل وتنفيذ العتاد.

يعرض الجدول عملية تنفيذ واحدة على الجهاز، كما في الملف الشخصي 2، لكن الشيفرة تسلك مسارًا مختلفًا. بدلاً من استدعاء عملية مدمجة مخصصة، يكتشف المترجم الدمج من دالة بايثون الأصلية. في هذا التشغيل، وقت Neuron قابل للمقارنة مع torch.rms_norm، عند 375 μs مقارنة بـ 403 μs.

ملخص جنباً إلى جنب

يدمج كلٌّ من torch.rms_norm و torch.compile عملية RMSNorm في إرسال نواة واحد. في هذه التجارب، يستغرق أيٌّ من النهجين حوالي سُدس زمن الجهاز للتنفيذ الساذج. يزيل الدمج الفجوات بين عمليات تنفيذ العتاد ويتجنب إنشاء قيم وسيطة مثل x_sq و mean_sq بين العمليات المُرسلة بشكل منفصل.

الخلاصة

key_averages() تُظهر أي العمليات تستهلك زمن الجهاز، بينما يكشف مخرجها الكامل أيضًا عن الإعداد والتحويلات على جانب المعالج. يوفر Explorer التتبع البصري المقابل. يمكنك متابعة عملية واحدة من إرسالها من إطار العمل مرورًا بجدولتها في بيئة التشغيل وصولًا إلى تنفيذها على العتاد، ومشاهدة الفجوات التي يزيلها الدمج.

يمكنك استخدام سير العمل نفسه على نماذج أكبر: ابحث عن نقطة ساخنة في الجدول، وتحقق منها في التتبع، ثم حدد ما إذا كانت عملية مدمجة أو torch.compile يمكن أن يخفض تكلفتها.

في الجزء 2، سننتقل إلى الملف الجهازي. ستفتح NEFF، وتفحص تفصيل TensorEngine وVectorEngine والوصول المباشر للذاكرة (DMA)، وتضع النواة على مخطط roofline لتحديد ما إذا كانت مقيدة بالحوسبة أم مقيدة بالذاكرة.

شكر خاص إلى Aritra Roy Gosthipaty على الإلهام لبدء هذه السلسلة، وإلى زملائي Pinak Panigrahi وSamir Araujo وAkhil Azhikodan وEsha Lakhotia والجميع على المراجعة!