การสร้างโปรไฟล์การคำนวณ JAX ด้วย XProf

XProf เป็นวิธีที่ยอดเยี่ยมในการรับและแสดงภาพร่องรอยและโปรไฟล์ประสิทธิภาพของโปรแกรม รวมถึงกิจกรรมใน GPU และ TPU ผลลัพธ์สุดท้ายจะมีลักษณะ ดังนี้

การจับภาพแบบเป็นโปรแกรม

คุณสามารถวัดโค้ดเพื่อบันทึกการติดตามโปรไฟล์สำหรับโค้ด JAX ผ่านเมธอด jax.profiler.start_trace และ jax.profiler.stop_trace ได้ Call jax.profiler.start_trace พร้อมไดเรกทอรีที่จะเขียนไฟล์การติดตาม ซึ่งควรเป็น--logdir ไดเรกทอรีเดียวกับที่ใช้เริ่ม XProf จากนั้นคุณจะใช้ XProf เพื่อดูการติดตามได้

เช่น หากต้องการบันทึกการติดตามของโปรไฟล์เลอร์ ให้ทำดังนี้

import jax

jax.profiler.start_trace("/tmp/profile-data")

# Run the operations to be profiled
key = jax.random.key(0)
x = jax.random.normal(key, (5000, 5000))
y = x @ x
y.block_until_ready()

jax.profiler.stop_trace()

โปรดทราบว่า jax.block_until_ready โทร เราใช้ข้อมูลนี้เพื่อให้มั่นใจว่าการดำเนินการในอุปกรณ์จะได้รับการบันทึกโดยการติดตาม ดูรายละเอียดเกี่ยวกับเหตุผลที่ต้องดำเนินการนี้ได้ที่การส่งแบบอะซิงโครนัส

คุณยังใช้เครื่องจัดการบริบท jax.profiler.trace แทน start_trace และ stop_trace ได้ด้วย

import jax

with jax.profiler.trace("/tmp/profile-data"):
  key = jax.random.key(0)
  x = jax.random.normal(key, (5000, 5000))
  y = x @ x
  y.block_until_ready()

การดูการติดตาม

หลังจากบันทึกการติดตามแล้ว คุณจะดูได้โดยใช้ UI ของ XProf

คุณเปิดใช้ UI ของโปรไฟล์เลอร์ได้โดยตรงโดยใช้คำสั่ง XProf แบบสแตนด์อโลนโดย ชี้ไปยังไดเรกทอรีบันทึก

$ xprof --port=8791 /tmp/profile-data
Attempting to start XProf server:
  Log Directory: /tmp/profile-data
  Port: 8791
  Worker Service Address: 0.0.0.0:50051
  Hide Capture Button: False
XProf at http://localhost:8791/ (Press CTRL+C to quit)

ไปที่ URL ที่ระบุ (เช่น http://localhost:8791/) ในเบราว์เซอร์ เพื่อดูโปรไฟล์

ร่องรอยที่ใช้ได้จะปรากฏในเมนูแบบเลื่อนลง "เซสชัน" ทางด้านซ้าย เลือกเซสชันที่สนใจ แล้วเลือก "Trace Viewer" ในเมนูแบบเลื่อนลง "เครื่องมือ" ตอนนี้คุณควรเห็นไทม์ไลน์ของการดำเนินการ คุณใช้ ปุ่ม WASD เพื่อไปยังส่วนต่างๆ ของการติดตาม และคลิกหรือลากเพื่อเลือกเหตุการณ์สำหรับ รายละเอียดเพิ่มเติมได้ ดูรายละเอียดเพิ่มเติมเกี่ยวกับการใช้เครื่องมือดูการติดตามได้ในเอกสารประกอบของเครื่องมือดูการติดตาม

การจับภาพด้วยตนเองผ่าน XProf

ต่อไปนี้เป็นวิธีการบันทึกการติดตาม N วินาทีที่ทริกเกอร์ด้วยตนเอง จากโปรแกรมที่กําลังทํางาน

1. เริ่มเซิร์ฟเวอร์ XProf

   xprof --logdir /tmp/profile-data/
   

   คุณควรโหลด XProf ได้ที่ <http://localhost:8791/> คุณสามารถระบุพอร์ตอื่นได้ด้วยแฟล็ก --port

2. ในโปรแกรมหรือกระบวนการ Python ที่ต้องการสร้างโปรไฟล์ ให้เพิ่มโค้ดต่อไปนี้ ที่ใดที่หนึ่งใกล้กับจุดเริ่มต้น

   import jax.profiler
   jax.profiler.start_server(9999)
   

   ซึ่งจะเริ่มเซิร์ฟเวอร์โปรไฟล์ที่ XProf เชื่อมต่อ เซิร์ฟเวอร์ Profiler ต้องทำงานก่อนที่คุณจะไปยังขั้นตอนถัดไป เมื่อใช้งานเซิร์ฟเวอร์เสร็จแล้ว คุณสามารถโทรหา jax.profiler.stop_server() เพื่อปิดเซิร์ฟเวอร์ได้

   หากต้องการสร้างโปรไฟล์ของข้อมูลโค้ดในโปรแกรมที่ทำงานเป็นเวลานาน (เช่น ลูปการฝึกที่ใช้เวลานาน ) คุณสามารถวางไว้ที่จุดเริ่มต้นของโปรแกรมและเริ่มโปรแกรมได้ตามปกติ หากต้องการสร้างโปรไฟล์โปรแกรมสั้นๆ (เช่น Microbenchmark) ตัวเลือกหนึ่งคือการเริ่มเซิร์ฟเวอร์โปรไฟล์ในเชลล์ IPython และเรียกใช้โปรแกรมสั้นๆ ด้วย %run หลังจากเริ่มการจับภาพใน ขั้นตอนถัดไป อีกทางเลือกหนึ่งคือการเริ่มเซิร์ฟเวอร์โปรไฟล์ที่ จุดเริ่มต้นของโปรแกรมและใช้ time.sleep() เพื่อให้มีเวลาเพียงพอในการ เริ่มการจับภาพ

3. เปิด <http://localhost:8791/> แล้วคลิกปุ่ม "CAPTURE PROFILE" ที่ด้านซ้ายบน ป้อน "localhost:9999" เป็น URL ของบริการโปรไฟล์ (นี่คือที่อยู่ของเซิร์ฟเวอร์โปรไฟล์ที่คุณเริ่มต้นในขั้นตอนก่อนหน้า) ป้อน จำนวนมิลลิวินาทีที่ต้องการสร้างโปรไฟล์ แล้วคลิก "จับภาพ"

4. หากโค้ดที่คุณต้องการสร้างโปรไฟล์ยังไม่ได้ทำงาน (เช่น หากคุณเริ่มเซิร์ฟเวอร์โปรไฟล์ในเชลล์ Python) ให้รันโค้ดขณะที่การจับภาพกำลังทำงาน

5. หลังจากจับภาพเสร็จแล้ว XProf ควรจะรีเฟรชโดยอัตโนมัติ (ฟีเจอร์การสร้างโปรไฟล์ XProf บางอย่างไม่ได้เชื่อมต่อกับ JAX ดังนั้นในตอนแรกจึงอาจดูเหมือนว่าไม่มีการบันทึกอะไรเลย) ทางด้านซ้ายในส่วน "เครื่องมือ" ให้เลือก "Trace Viewer"

ตอนนี้คุณควรเห็นไทม์ไลน์ของการดำเนินการ คุณสามารถใช้ปุ่ม WASD เพื่อ ไปยังส่วนต่างๆ ของการติดตาม และคลิกหรือลากเพื่อเลือกเหตุการณ์เพื่อดูรายละเอียดเพิ่มเติมที่ ด้านล่าง ดูรายละเอียดเพิ่มเติมเกี่ยวกับการใช้เครื่องมือ Trace Viewer ได้ที่เอกสารประกอบของเครื่องมือ Trace Viewer

XProf และ Tensorboard

XProf เป็นเครื่องมือพื้นฐานที่ขับเคลื่อนฟังก์ชันการทำงานของการสร้างโปรไฟล์และการบันทึกการติดตามใน Tensorboard ตราบใดที่ติดตั้ง xprof แล้ว แท็บ "โปรไฟล์" จะอยู่ใน Tensorboard การใช้ฟีเจอร์นี้จะเหมือนกับการเปิดใช้ XProf แยกกัน ตราบใดที่เปิดใช้โดยชี้ไปยังไดเรกทอรีบันทึกเดียวกัน ซึ่งรวมถึงฟังก์ชันการจับภาพ การวิเคราะห์ และการดูโปรไฟล์ XProf จะแทนที่tensorboard_plugin_profileฟังก์ชันการทำงานที่ก่อนหน้านี้เรา แนะนำ

$ tensorboard --logdir=/tmp/profile-data
[...]
Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all
TensorBoard 2.19.0 at http://localhost:6006/ (Press CTRL+C to quit)

การเพิ่มเหตุการณ์การติดตามที่กำหนดเอง

โดยค่าเริ่มต้น เหตุการณ์ในเครื่องมือแสดงข้อมูลการติดตามส่วนใหญ่จะเป็นฟังก์ชัน JAX ภายในระดับต่ำ คุณเพิ่มเหตุการณ์และฟังก์ชันของคุณเองได้โดยใช้ jax.profiler.TraceAnnotation และ jax.profiler.annotate_function ในโค้ด

การกำหนดค่าตัวเลือกโปรไฟล์

เมธอด start_trace รับพารามิเตอร์ profiler_options ที่ไม่บังคับ ซึ่ง ช่วยให้ควบคุมลักษณะการทำงานของโปรไฟล์เลอร์ได้อย่างละเอียด พารามิเตอร์นี้ ควรเป็นอินสแตนซ์ของ jax.profiler.ProfileOptions

เช่น หากต้องการปิดใช้การติดตามทั้งหมดของ Python และโฮสต์ ให้ทำดังนี้

import jax

options = jax.profiler.ProfileOptions()
options.python_tracer_level = 0
options.host_tracer_level = 0
jax.profiler.start_trace("/tmp/profile-data", profiler_options=options)

# Run the operations to be profiled
key = jax.random.key(0)
x = jax.random.normal(key, (5000, 5000))
y = x @ x
y.block_until_ready()

jax.profiler.stop_trace()

ตัวเลือกทั่วไป

1. host_tracer_level: ตั้งค่าระดับการติดตามสำหรับกิจกรรมฝั่งโฮสต์

   ค่าที่รองรับ

   • 0: ปิดใช้การติดตามโฮสต์ (CPU) ทั้งหมด
   • 1: เปิดใช้การติดตามเฉพาะเหตุการณ์ TraceMe ที่ผู้ใช้สร้างเครื่องมือ
   • 2: มีการติดตามระดับ 1 พร้อมรายละเอียดการดำเนินการโปรแกรมระดับสูง เช่น การดำเนินการ XLA ที่มีค่าใช้จ่ายสูง (ค่าเริ่มต้น)
   • 3: มีการติดตามระดับ 2 พร้อมรายละเอียดการดำเนินการโปรแกรมระดับต่ำที่ละเอียดมากขึ้น เช่น การดำเนินการ XLA ที่มีราคาถูก

2. device_tracer_level: ควบคุมว่าจะเปิดใช้การติดตามอุปกรณ์หรือไม่

   ค่าที่รองรับ

   • 0: ปิดใช้การติดตามอุปกรณ์
   • 1: เปิดใช้การติดตามอุปกรณ์ (ค่าเริ่มต้น)

3. python_tracer_level: ควบคุมว่าจะเปิดใช้การติดตาม Python หรือไม่

   ค่าที่รองรับ

   • 0: ปิดใช้การติดตามการเรียกฟังก์ชัน Python (ค่าเริ่มต้น)
   • 1: เปิดใช้การติดตาม Python

ตัวเลือกการกำหนดค่าขั้นสูง

ตัวเลือก TPU

1. tpu_trace_mode: ระบุโหมดสำหรับการติดตาม TPU

   ค่าที่รองรับ

   • TRACE_ONLY_HOST: ซึ่งหมายความว่าจะมีการติดตามเฉพาะกิจกรรมฝั่งโฮสต์ (CPU) และจะไม่มีการรวบรวมการติดตามอุปกรณ์ (TPU/GPU)
   • TRACE_ONLY_XLA: ซึ่งหมายความว่าจะมีการติดตามเฉพาะการดำเนินการระดับ XLA ในอุปกรณ์
   • TRACE_COMPUTE: ติดตามการดำเนินการประมวลผลในอุปกรณ์
   • TRACE_COMPUTE_AND_SYNC: ติดตามทั้งการดำเนินการประมวลผลและ เหตุการณ์การซิงค์ในอุปกรณ์

   หากไม่ได้ระบุ "tpu_trace_mode" ค่าเริ่มต้นของ trace_mode จะเป็น TRACE_ONLY_XLA

2. tpu_num_sparse_cores_to_trace: ระบุจำนวน Sparse Core ที่จะ ติดตามใน TPU

3. tpu_num_sparse_core_tiles_to_trace: ระบุจำนวนไทล์ภายใน แต่ละคอร์แบบกระจัดกระจายที่จะติดตามใน TPU

4. tpu_num_chips_to_profile_per_task: ระบุจำนวนชิป TPU ที่จะ สร้างโปรไฟล์ต่อชิ้นงาน

ตัวเลือก GPU

ตัวเลือกต่อไปนี้ใช้ได้สำหรับการสร้างโปรไฟล์ GPU

• gpu_max_callback_api_events: กำหนดจำนวนเหตุการณ์สูงสุดที่รวบรวมโดย API การเรียกกลับของ CUPTI ค่าเริ่มต้นคือ 2*1024*1024
• gpu_max_activity_api_events: กำหนดจำนวนเหตุการณ์สูงสุดที่รวบรวมโดย API กิจกรรม CUPTI ค่าเริ่มต้นคือ 2*1024*1024
• gpu_max_annotation_strings: กำหนดจำนวนสูงสุดของสตริงคำอธิบายประกอบ ที่รวบรวมได้ ค่าเริ่มต้นคือ 1024*1024
• gpu_enable_nvtx_tracking: เปิดใช้การติดตาม NVTX ใน CUPTI ค่าเริ่มต้นคือ False
• gpu_enable_cupti_activity_graph_trace: เปิดใช้การติดตามกราฟกิจกรรม CUPTI สำหรับกราฟ CUDA ค่าเริ่มต้นคือ False
• gpu_pm_sample_counters: สตริงที่คั่นด้วยคอมมาของเมตริกการตรวจสอบประสิทธิภาพ GPU เพื่อรวบรวมโดยใช้ฟีเจอร์การสุ่มตัวอย่าง PM ของ CUPTI (เช่น "sm__cycles_active.avg.pct_of_peak_sustained_elapsed") การสุ่มตัวอย่าง PM จะ ปิดใช้โดยค่าเริ่มต้น ดูเมตริกที่ใช้ได้ในเอกสารประกอบ CUPTI ของ NVIDIA
• gpu_pm_sample_interval_us: กำหนดช่วงเวลาการสุ่มตัวอย่างในหน่วยไมโครวินาทีสำหรับ การสุ่มตัวอย่าง PM ของ CUPTI ค่าเริ่มต้นคือ 500
• gpu_pm_sample_buffer_size_per_gpu_mb: กำหนดขนาดบัฟเฟอร์หน่วยความจำของระบบต่อ อุปกรณ์เป็น MB สำหรับการสุ่มตัวอย่าง PM ของ CUPTI ค่าเริ่มต้นคือ 64 MB ค่าสูงสุดที่รองรับคือ 4 GB
• gpu_num_chips_to_profile_per_task: ระบุจำนวนอุปกรณ์ GPU ที่จะ สร้างโปรไฟล์ต่อ 1 งาน หากไม่ได้ระบุ ตั้งค่าเป็น 0 หรือตั้งค่าเป็นค่าที่ไม่ถูกต้อง ระบบจะสร้างโปรไฟล์ GPU ที่พร้อมใช้งานทั้งหมด ซึ่งใช้เพื่อลดขนาดการรวบรวม การติดตามได้
• gpu_dump_graph_node_mapping: หากเปิดใช้ จะทิ้งข้อมูลการแมปโหนดกราฟ CUDA ลงในการติดตาม ค่าเริ่มต้นคือ False

เช่น

options = ProfileOptions()
options.advanced_configuration = {"tpu_trace_mode" : "TRACE_ONLY_HOST", "tpu_num_sparse_cores_to_trace" : 2}

แสดงผล InvalidArgumentError หากพบคีย์หรือค่าตัวเลือกที่ไม่รู้จัก