|
25 | 25 | pip install -U "jax[cuda13]" |
26 | 26 | pip install numpyro |
27 | 27 | python scripts/test-jax-install.py |
28 | | - # === lax.scan GPU Performance Profiling === |
29 | | - - name: Profile lax.scan (GPU vs CPU) |
30 | | - shell: bash -l {0} |
31 | | - run: | |
32 | | - echo "=== lax.scan Performance Profiling ===" |
33 | | - echo "This profiles the known issue with lax.scan on GPU (JAX Issue #2491)" |
34 | | - echo "" |
35 | | - python scripts/profile_lax_scan.py --iterations 100000 --diagnose |
36 | | - echo "" |
37 | | - echo "The diagnostic shows if time scales linearly with iterations," |
38 | | - echo "which indicates constant per-iteration CPU-GPU sync overhead." |
39 | | - - name: Nsight Systems Profile (if available) |
40 | | - shell: bash -l {0} |
41 | | - continue-on-error: true |
42 | | - run: | |
43 | | - echo "=== NVIDIA Nsight Systems Profiling ===" |
44 | | - if command -v nsys &> /dev/null; then |
45 | | - echo "nsys found, running profile with 1000 iterations..." |
46 | | - mkdir -p nsight_profiles |
47 | | - nsys profile -o nsight_profiles/lax_scan_trace \ |
48 | | - --trace=cuda,nvtx,osrt \ |
49 | | - --cuda-memory-usage=true \ |
50 | | - --stats=true \ |
51 | | - python scripts/profile_lax_scan.py --nsys -n 1000 |
52 | | - echo "" |
53 | | - echo "Profile saved to nsight_profiles/lax_scan_trace.nsys-rep" |
54 | | - echo "Download artifact and open in Nsight Systems UI to see CPU-GPU sync pattern" |
55 | | - else |
56 | | - echo "nsys not found, skipping Nsight profiling" |
57 | | - echo "Install NVIDIA Nsight Systems to enable this profiling" |
58 | | - fi |
59 | | - - name: Upload Nsight Profile |
60 | | - uses: actions/upload-artifact@v5 |
61 | | - if: success() || failure() |
62 | | - continue-on-error: true |
63 | | - with: |
64 | | - name: nsight-profile |
65 | | - path: nsight_profiles/ |
66 | | - if-no-files-found: ignore |
67 | | - # === Benchmark Tests (Bare Metal, Jupyter, Jupyter-Book) === |
68 | | - - name: Run Hardware Benchmarks (Bare Metal) |
69 | | - shell: bash -l {0} |
70 | | - run: | |
71 | | - echo "=== Bare Metal Python Script Execution ===" |
72 | | - python scripts/benchmark-hardware.py |
73 | | - mkdir -p benchmark_results |
74 | | - mv benchmark_results_bare_metal.json benchmark_results/ |
75 | | - - name: Run Jupyter Notebook Benchmark (via nbconvert) |
76 | | - shell: bash -l {0} |
77 | | - run: | |
78 | | - echo "=== Jupyter Kernel Execution ===" |
79 | | - cd scripts |
80 | | - jupyter nbconvert --to notebook --execute benchmark-jupyter.ipynb --output benchmark-jupyter-executed.ipynb |
81 | | - echo "Notebook executed successfully" |
82 | | - cd .. |
83 | | - mv scripts/benchmark_results_jupyter.json benchmark_results/ |
84 | | - - name: Run Jupyter-Book Benchmark |
85 | | - shell: bash -l {0} |
86 | | - run: | |
87 | | - echo "=== Jupyter-Book Execution ===" |
88 | | - # Build just the benchmark file using jupyter-book |
89 | | - mkdir -p benchmark_test |
90 | | - cp scripts/benchmark-jupyterbook.md benchmark_test/ |
91 | | - # Create minimal _config.yml |
92 | | - echo "title: Benchmark Test" > benchmark_test/_config.yml |
93 | | - echo "execute:" >> benchmark_test/_config.yml |
94 | | - echo " execute_notebooks: force" >> benchmark_test/_config.yml |
95 | | - # Create minimal _toc.yml |
96 | | - echo "format: jb-book" > benchmark_test/_toc.yml |
97 | | - echo "root: benchmark-jupyterbook" >> benchmark_test/_toc.yml |
98 | | - # Build (run from benchmark_test so JSON is written there) |
99 | | - cd benchmark_test |
100 | | - jb build . --path-output ../benchmark_build/ |
101 | | - cd .. |
102 | | - echo "Jupyter-Book build completed successfully" |
103 | | - # Move JSON results if generated |
104 | | - cp benchmark_test/benchmark_results_jupyterbook.json benchmark_results/ 2>/dev/null || echo "No jupyterbook results" |
105 | | - - name: Collect and Display Benchmark Results |
106 | | - shell: bash -l {0} |
107 | | - run: | |
108 | | - echo "=== Benchmark Results Summary ===" |
109 | | - for f in benchmark_results/*.json; do |
110 | | - echo "--- $f ---" |
111 | | - cat "$f" |
112 | | - echo "" |
113 | | - done |
114 | | - - name: Upload Benchmark Results |
115 | | - uses: actions/upload-artifact@v5 |
116 | | - with: |
117 | | - name: benchmark-results |
118 | | - path: benchmark_results/ |
119 | | - if-no-files-found: warn |
120 | 28 | - name: Install latex dependencies |
121 | 29 | run: | |
122 | 30 | sudo apt-get -qq update |
|
0 commit comments