diff --git a/docs/profiling/budget_formula_sweep.json b/docs/profiling/budget_formula_sweep.json new file mode 100644 index 0000000000..22e5fa252b --- /dev/null +++ b/docs/profiling/budget_formula_sweep.json @@ -0,0 +1,44 @@ +[ + { + "vocab_size": 8000, + "auto_chunk_size": 524, + "single_chunk_peak_bytes": 50332672, + "full_run_peak_extra_bytes": 264242176, + "budget_bytes": 33554432 + }, + { + "vocab_size": 16000, + "auto_chunk_size": 262, + "single_chunk_peak_bytes": 50332672, + "full_run_peak_extra_bytes": 524427264, + "budget_bytes": 33554432 + }, + { + "vocab_size": 32000, + "auto_chunk_size": 131, + "single_chunk_peak_bytes": 50332672, + "full_run_peak_extra_bytes": 1048862720, + "budget_bytes": 33554432 + }, + { + "vocab_size": 50304, + "auto_chunk_size": 83, + "single_chunk_peak_bytes": 50332672, + "full_run_peak_extra_bytes": 1652100608, + "budget_bytes": 33554432 + }, + { + "vocab_size": 100000, + "auto_chunk_size": 41, + "single_chunk_peak_bytes": 50332672, + "full_run_peak_extra_bytes": 3314145792, + "budget_bytes": 33554432 + }, + { + "vocab_size": 152064, + "auto_chunk_size": 27, + "single_chunk_peak_bytes": 50332672, + "full_run_peak_extra_bytes": 5037360128, + "budget_bytes": 33554432 + } +] diff --git a/docs/profiling/budget_formula_sweep.png b/docs/profiling/budget_formula_sweep.png new file mode 100644 index 0000000000..b93fb6dbec Binary files /dev/null and b/docs/profiling/budget_formula_sweep.png differ diff --git a/docs/profiling/compile_vs_eager.json b/docs/profiling/compile_vs_eager.json new file mode 100644 index 0000000000..3939bd2594 --- /dev/null +++ b/docs/profiling/compile_vs_eager.json @@ -0,0 +1,22 @@ +{ + "chunk_0": { + "eager": { + "peak_bytes": 2097185792, + "elapsed_ms": 160.83151245117188 + }, + "compiled": { + "peak_bytes": 1048643072, + "elapsed_ms": 6.78707218170166 + } + }, + "chunk_128": { + "eager": { + "peak_bytes": 1585480704, + "elapsed_ms": 222.7342987060547 + }, + "compiled": { + "peak_bytes": 1065419264, + "elapsed_ms": 7.819295883178711 + } + } +} diff --git a/docs/profiling/compile_vs_eager.png b/docs/profiling/compile_vs_eager.png new file mode 100644 index 0000000000..3c32629d8e Binary files /dev/null and b/docs/profiling/compile_vs_eager.png differ diff --git a/docs/profiling/end_to_end_step.json b/docs/profiling/end_to_end_step.json new file mode 100644 index 0000000000..3f430dc2db --- /dev/null +++ b/docs/profiling/end_to_end_step.json @@ -0,0 +1,10 @@ +{ + "model_name": "pythia-14m", + "vocab_size": 50304, + "seq_length": 512, + "batch_size": 8, + "peak_unchunked_bytes": 3580535808, + "peak_chunked_bytes": 2764874752, + "table_unchunked": "------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ \n Name Self CPU % Self CPU CPU total % CPU total CPU time avg Self CUDA Self CUDA % CUDA total CUDA time avg CPU Mem Self CPU Mem CUDA Mem Self CUDA Mem # of Calls \n------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ \n aten::mm 4.59% 33.924ms 8.62% 63.641ms 1.248ms 60.906ms 49.38% 79.954ms 1.568ms 0 B 0 B 909.19 MB 909.19 MB 51 \n aten::_log_softmax 3.41% 25.159ms 6.34% 46.822ms 46.822ms 12.868ms 10.43% 25.736ms 25.736ms 0 B 0 B 786.00 MB 786.00 MB 1 \n aten::nll_loss_backward 0.37% 2.706ms 0.47% 3.476ms 3.476ms 25.948us 0.02% 3.587ms 3.587ms 0 B 0 B 786.00 MB 786.00 MB 1 \n aten::_log_softmax_backward_data 0.01% 71.293us 0.08% 626.965us 626.965us 13.043ms 10.57% 13.043ms 13.043ms 0 B 0 B 786.00 MB 786.00 MB 1 \n aten::empty 0.30% 2.194ms 0.91% 6.718ms 32.933us 0.000us 0.00% 0.000us 0.000us 96 B 96 B 236.15 MB 236.15 MB 204 \n aten::addmm 0.95% 7.013ms 8.23% 60.752ms 2.531ms 4.617ms 3.74% 5.130ms 213.756us 0 B 0 B 116.12 MB 116.12 MB 24 \n aten::cat 4.85% 35.799ms 5.96% 44.042ms 1.468ms 857.023us 0.69% 870.941us 29.031us 0 B 0 B 66.00 MB 66.00 MB 30 \n aten::gelu 4.76% 35.142ms 8.35% 61.671ms 10.279ms 433.023us 0.35% 505.492us 84.249us 0 B 0 B 48.00 MB 48.00 MB 6 \n aten::gelu_backward 0.02% 159.754us 0.06% 442.076us 73.679us 622.145us 0.50% 622.145us 103.691us 0 B 0 B 48.00 MB 48.00 MB 6 \n aten::add 2.06% 15.207ms 3.68% 27.144ms 904.804us 499.731us 0.41% 535.918us 17.864us 0 B 0 B 42.00 MB 42.00 MB 30 \n aten::mul 1.44% 10.601ms 2.53% 18.685ms 346.027us 624.287us 0.51% 633.821us 11.737us 0 B 0 B 36.00 MB 36.00 MB 54 \n aten::neg 1.94% 14.311ms 3.27% 24.137ms 1.006ms 200.932us 0.16% 212.802us 8.867us 0 B 0 B 6.00 MB 6.00 MB 24 \n aten::resize_ 0.01% 95.246us 0.39% 2.910ms 970.154us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 2.03 MB 2.03 MB 3 \n aten::sum 2.51% 18.552ms 5.71% 42.181ms 1.406ms 1.448ms 1.17% 1.448ms 48.273us 0 B 0 B 411.00 KB 411.00 KB 30 \n aten::nll_loss_forward 1.01% 7.480ms 1.76% 12.986ms 12.986ms 31.963us 0.03% 63.926us 63.926us 0 B 0 B 1.00 KB 1.00 KB 1 \n aten::empty_strided 0.00% 24.476us 0.00% 24.476us 24.476us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 512 B 512 B 1 \n aten::slice 0.09% 632.750us 0.11% 793.038us 5.432us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 146 \n aten::as_strided 0.08% 574.081us 0.08% 574.081us 1.348us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 426 \n aten::unsqueeze 0.00% 17.726us 0.00% 19.395us 9.697us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 2 \n aten::embedding 0.01% 90.507us 6.01% 44.382ms 44.382ms 0.000us 0.00% 39.802us 39.802us 0 B 0 B 2.00 MB 0 B 1 \n------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ \nSelf CPU time total: 738.597ms\nSelf CUDA time total: 123.343ms\n", + "table_chunked": "------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ \n Name Self CPU % Self CPU CPU total % CPU total CPU time avg Self CUDA Self CUDA % CUDA total CUDA time avg CPU Mem Self CPU Mem CUDA Mem Self CUDA Mem # of Calls \n------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ \n aten::mm 0.70% 1.291ms 1.57% 2.902ms 56.905us 65.031ms 45.73% 65.031ms 1.275ms 0 B 0 B 901.06 MB 901.06 MB 51 \n aten::cat 0.36% 658.672us 1.06% 1.963ms 61.359us 7.856ms 5.52% 7.861ms 245.643us 0 B 0 B 852.02 MB 852.02 MB 32 \n aten::_log_softmax 0.46% 855.081us 3.68% 6.815ms 212.978us 16.445ms 11.56% 16.445ms 513.894us 0 B 0 B 786.00 MB 786.00 MB 32 \n aten::_log_softmax_backward_data 0.21% 381.745us 0.50% 928.148us 29.005us 15.293ms 10.75% 15.293ms 477.894us 0 B 0 B 786.00 MB 786.00 MB 32 \n aten::nll_loss_backward 0.28% 510.827us 2.04% 3.773ms 117.904us 151.103us 0.11% 4.188ms 130.874us 0 B 0 B 786.00 MB 785.98 MB 32 \n aten::empty 0.76% 1.400ms 1.05% 1.944ms 8.239us 0.000us 0.00% 0.000us 0.000us 96 B 96 B 236.16 MB 236.16 MB 236 \n aten::addmm 0.52% 966.759us 1.22% 2.264ms 94.354us 4.957ms 3.49% 4.957ms 206.534us 0 B 0 B 108.00 MB 108.00 MB 24 \n aten::gelu 0.07% 133.055us 0.34% 626.038us 104.340us 485.558us 0.34% 485.558us 80.926us 0 B 0 B 48.00 MB 48.00 MB 6 \n aten::gelu_backward 0.05% 87.548us 0.08% 143.945us 23.991us 620.754us 0.44% 620.754us 103.459us 0 B 0 B 48.00 MB 48.00 MB 6 \n aten::add 0.29% 538.806us 1.30% 2.401ms 80.036us 522.004us 0.37% 522.004us 17.400us 0 B 0 B 42.00 MB 42.00 MB 30 \n aten::mul 0.45% 824.568us 0.83% 1.543ms 28.579us 644.691us 0.45% 644.691us 11.939us 0 B 0 B 36.00 MB 36.00 MB 54 \n aten::neg 0.21% 388.406us 0.33% 613.127us 25.547us 204.636us 0.14% 204.636us 8.527us 0 B 0 B 6.00 MB 6.00 MB 24 \n aten::resize_ 0.01% 19.952us 0.01% 19.952us 9.976us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 2.03 MB 2.03 MB 2 \n aten::sum 0.40% 746.197us 0.72% 1.340ms 41.885us 1.483ms 1.04% 1.510ms 47.181us 0 B 0 B 412.00 KB 380.00 KB 32 \n aten::empty_strided 0.02% 37.059us 0.02% 37.059us 12.353us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 33.00 KB 33.00 KB 3 \n aten::nll_loss_forward 0.33% 615.660us 0.78% 1.449ms 45.292us 158.556us 0.11% 266.937us 8.342us 0 B 0 B 32.00 KB 32.00 KB 32 \n aten::ne 3.16% 5.845ms 15.84% 29.324ms 29.324ms 4.448us 0.00% 8.896us 8.896us 0 B 0 B 4.00 KB 4.00 KB 1 \n aten::div 6.30% 11.668ms 10.70% 19.794ms 9.897ms 9.825us 0.01% 14.977us 7.488us 0 B 0 B 1.00 KB 1.00 KB 2 \n aten::maximum 5.99% 11.084ms 12.36% 22.866ms 22.866ms 3.808us 0.00% 7.616us 7.616us 0 B 0 B 512 B 512 B 1 \n aten::slice 0.36% 663.813us 0.48% 893.968us 3.694us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 242 \n------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ \nSelf CPU time total: 185.071ms\nSelf CUDA time total: 142.206ms\n" +} diff --git a/docs/profiling/end_to_end_step.png b/docs/profiling/end_to_end_step.png new file mode 100644 index 0000000000..c47f8095b5 Binary files /dev/null and b/docs/profiling/end_to_end_step.png differ diff --git a/docs/profiling/end_to_end_step_op_table.md b/docs/profiling/end_to_end_step_op_table.md new file mode 100644 index 0000000000..fc6b886577 --- /dev/null +++ b/docs/profiling/end_to_end_step_op_table.md @@ -0,0 +1,61 @@ +### unchunked (chunk_size=0) + +``` +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + Name Self CPU % Self CPU CPU total % CPU total CPU time avg Self CUDA Self CUDA % CUDA total CUDA time avg CPU Mem Self CPU Mem CUDA Mem Self CUDA Mem # of Calls +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + aten::mm 4.59% 33.924ms 8.62% 63.641ms 1.248ms 60.906ms 49.38% 79.954ms 1.568ms 0 B 0 B 909.19 MB 909.19 MB 51 + aten::_log_softmax 3.41% 25.159ms 6.34% 46.822ms 46.822ms 12.868ms 10.43% 25.736ms 25.736ms 0 B 0 B 786.00 MB 786.00 MB 1 + aten::nll_loss_backward 0.37% 2.706ms 0.47% 3.476ms 3.476ms 25.948us 0.02% 3.587ms 3.587ms 0 B 0 B 786.00 MB 786.00 MB 1 + aten::_log_softmax_backward_data 0.01% 71.293us 0.08% 626.965us 626.965us 13.043ms 10.57% 13.043ms 13.043ms 0 B 0 B 786.00 MB 786.00 MB 1 + aten::empty 0.30% 2.194ms 0.91% 6.718ms 32.933us 0.000us 0.00% 0.000us 0.000us 96 B 96 B 236.15 MB 236.15 MB 204 + aten::addmm 0.95% 7.013ms 8.23% 60.752ms 2.531ms 4.617ms 3.74% 5.130ms 213.756us 0 B 0 B 116.12 MB 116.12 MB 24 + aten::cat 4.85% 35.799ms 5.96% 44.042ms 1.468ms 857.023us 0.69% 870.941us 29.031us 0 B 0 B 66.00 MB 66.00 MB 30 + aten::gelu 4.76% 35.142ms 8.35% 61.671ms 10.279ms 433.023us 0.35% 505.492us 84.249us 0 B 0 B 48.00 MB 48.00 MB 6 + aten::gelu_backward 0.02% 159.754us 0.06% 442.076us 73.679us 622.145us 0.50% 622.145us 103.691us 0 B 0 B 48.00 MB 48.00 MB 6 + aten::add 2.06% 15.207ms 3.68% 27.144ms 904.804us 499.731us 0.41% 535.918us 17.864us 0 B 0 B 42.00 MB 42.00 MB 30 + aten::mul 1.44% 10.601ms 2.53% 18.685ms 346.027us 624.287us 0.51% 633.821us 11.737us 0 B 0 B 36.00 MB 36.00 MB 54 + aten::neg 1.94% 14.311ms 3.27% 24.137ms 1.006ms 200.932us 0.16% 212.802us 8.867us 0 B 0 B 6.00 MB 6.00 MB 24 + aten::resize_ 0.01% 95.246us 0.39% 2.910ms 970.154us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 2.03 MB 2.03 MB 3 + aten::sum 2.51% 18.552ms 5.71% 42.181ms 1.406ms 1.448ms 1.17% 1.448ms 48.273us 0 B 0 B 411.00 KB 411.00 KB 30 + aten::nll_loss_forward 1.01% 7.480ms 1.76% 12.986ms 12.986ms 31.963us 0.03% 63.926us 63.926us 0 B 0 B 1.00 KB 1.00 KB 1 + aten::empty_strided 0.00% 24.476us 0.00% 24.476us 24.476us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 512 B 512 B 1 + aten::slice 0.09% 632.750us 0.11% 793.038us 5.432us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 146 + aten::as_strided 0.08% 574.081us 0.08% 574.081us 1.348us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 426 + aten::unsqueeze 0.00% 17.726us 0.00% 19.395us 9.697us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 2 + aten::embedding 0.01% 90.507us 6.01% 44.382ms 44.382ms 0.000us 0.00% 39.802us 39.802us 0 B 0 B 2.00 MB 0 B 1 +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ +Self CPU time total: 738.597ms +Self CUDA time total: 123.343ms +``` + +### chunked (chunk_size=128) + +``` +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + Name Self CPU % Self CPU CPU total % CPU total CPU time avg Self CUDA Self CUDA % CUDA total CUDA time avg CPU Mem Self CPU Mem CUDA Mem Self CUDA Mem # of Calls +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + aten::mm 0.70% 1.291ms 1.57% 2.902ms 56.905us 65.031ms 45.73% 65.031ms 1.275ms 0 B 0 B 901.06 MB 901.06 MB 51 + aten::cat 0.36% 658.672us 1.06% 1.963ms 61.359us 7.856ms 5.52% 7.861ms 245.643us 0 B 0 B 852.02 MB 852.02 MB 32 + aten::_log_softmax 0.46% 855.081us 3.68% 6.815ms 212.978us 16.445ms 11.56% 16.445ms 513.894us 0 B 0 B 786.00 MB 786.00 MB 32 + aten::_log_softmax_backward_data 0.21% 381.745us 0.50% 928.148us 29.005us 15.293ms 10.75% 15.293ms 477.894us 0 B 0 B 786.00 MB 786.00 MB 32 + aten::nll_loss_backward 0.28% 510.827us 2.04% 3.773ms 117.904us 151.103us 0.11% 4.188ms 130.874us 0 B 0 B 786.00 MB 785.98 MB 32 + aten::empty 0.76% 1.400ms 1.05% 1.944ms 8.239us 0.000us 0.00% 0.000us 0.000us 96 B 96 B 236.16 MB 236.16 MB 236 + aten::addmm 0.52% 966.759us 1.22% 2.264ms 94.354us 4.957ms 3.49% 4.957ms 206.534us 0 B 0 B 108.00 MB 108.00 MB 24 + aten::gelu 0.07% 133.055us 0.34% 626.038us 104.340us 485.558us 0.34% 485.558us 80.926us 0 B 0 B 48.00 MB 48.00 MB 6 + aten::gelu_backward 0.05% 87.548us 0.08% 143.945us 23.991us 620.754us 0.44% 620.754us 103.459us 0 B 0 B 48.00 MB 48.00 MB 6 + aten::add 0.29% 538.806us 1.30% 2.401ms 80.036us 522.004us 0.37% 522.004us 17.400us 0 B 0 B 42.00 MB 42.00 MB 30 + aten::mul 0.45% 824.568us 0.83% 1.543ms 28.579us 644.691us 0.45% 644.691us 11.939us 0 B 0 B 36.00 MB 36.00 MB 54 + aten::neg 0.21% 388.406us 0.33% 613.127us 25.547us 204.636us 0.14% 204.636us 8.527us 0 B 0 B 6.00 MB 6.00 MB 24 + aten::resize_ 0.01% 19.952us 0.01% 19.952us 9.976us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 2.03 MB 2.03 MB 2 + aten::sum 0.40% 746.197us 0.72% 1.340ms 41.885us 1.483ms 1.04% 1.510ms 47.181us 0 B 0 B 412.00 KB 380.00 KB 32 + aten::empty_strided 0.02% 37.059us 0.02% 37.059us 12.353us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 33.00 KB 33.00 KB 3 + aten::nll_loss_forward 0.33% 615.660us 0.78% 1.449ms 45.292us 158.556us 0.11% 266.937us 8.342us 0 B 0 B 32.00 KB 32.00 KB 32 + aten::ne 3.16% 5.845ms 15.84% 29.324ms 29.324ms 4.448us 0.00% 8.896us 8.896us 0 B 0 B 4.00 KB 4.00 KB 1 + aten::div 6.30% 11.668ms 10.70% 19.794ms 9.897ms 9.825us 0.01% 14.977us 7.488us 0 B 0 B 1.00 KB 1.00 KB 2 + aten::maximum 5.99% 11.084ms 12.36% 22.866ms 22.866ms 3.808us 0.00% 7.616us 7.616us 0 B 0 B 512 B 512 B 1 + aten::slice 0.36% 663.813us 0.48% 893.968us 3.694us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 242 +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ +Self CPU time total: 185.071ms +Self CUDA time total: 142.206ms +``` diff --git a/docs/profiling/memory_timeline_chunk128.png b/docs/profiling/memory_timeline_chunk128.png new file mode 100644 index 0000000000..66456f2a34 Binary files /dev/null and b/docs/profiling/memory_timeline_chunk128.png differ diff --git a/docs/profiling/memory_timeline_unchunked.png b/docs/profiling/memory_timeline_unchunked.png new file mode 100644 index 0000000000..32093a52a4 Binary files /dev/null and b/docs/profiling/memory_timeline_unchunked.png differ diff --git a/docs/profiling/op_table.md b/docs/profiling/op_table.md new file mode 100644 index 0000000000..ab8771b254 --- /dev/null +++ b/docs/profiling/op_table.md @@ -0,0 +1,155 @@ +### chunk_size=0 + +``` +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + Name Self CPU % Self CPU CPU total % CPU total CPU time avg CPU Mem Self CPU Mem # of Calls +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + aten::_log_softmax 7.35% 32.840ms 7.35% 32.840ms 32.840ms 500.00 MB 500.00 MB 1 + aten::nll_loss_backward 0.41% 1.824ms 20.28% 90.546ms 90.546ms 500.00 MB 500.00 MB 1 + aten::_log_softmax_backward_data 68.27% 304.859ms 68.27% 304.859ms 304.859ms 500.00 MB 500.00 MB 1 + aten::nll_loss_forward 0.25% 1.101ms 0.25% 1.101ms 1.101ms 8 B 8 B 1 + aten::empty_strided 0.00% 2.292us 0.00% 2.292us 2.292us 4 B 4 B 1 + (198): _run_module_as_main 0.00% 0.917us 100.00% 446.555ms 446.555ms 500.00 MB 0 B 1 + (88): _run_code 0.00% 1.458us 100.00% 446.555ms 446.555ms 500.00 MB 0 B 1 + litgpt/scripts/profile_memory.py(125): 0.00% 1.125us 100.00% 446.553ms 446.553ms 500.00 MB 0 B 1 + litgpt/scripts/profile_memory.py(70): main 0.00% 1.333us 100.00% 446.552ms 446.552ms 500.00 MB 0 B 1 +litgpt/scripts/profile_memory.py(33): profile_chunk_... 0.01% 42.680us 100.00% 446.551ms 446.551ms 500.00 MB 0 B 1 + torch/profiler/profiler.py(1118): __enter__ 0.00% 0.416us 0.01% 65.118us 65.118us 0 B 0 B 1 + torch/profiler/profiler.py(1128): start 0.00% 0.459us 0.01% 64.702us 64.702us 0 B 0 B 1 + torch/profiler/profiler.py(1262): _transit_action 0.00% 1.792us 0.01% 64.243us 64.243us 0 B 0 B 1 + torch/profiler/profiler.py(342): start_trace 0.00% 13.082us 0.01% 62.451us 62.451us 0 B 0 B 1 + torch/autograd/profiler.py(414): _start_trace 0.01% 29.392us 0.01% 30.892us 30.892us 0 B 0 B 1 +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ +Self CPU time total: 446.555ms + +``` + +### chunk_size=32 + +``` +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + Name Self CPU % Self CPU CPU total % CPU total CPU time avg CPU Mem Self CPU Mem # of Calls +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + aten::cat 40.86% 102.629ms 40.87% 102.635ms 51.318ms 500.02 MB 500.02 MB 2 + aten::_log_softmax 31.04% 77.958ms 31.04% 77.958ms 609.050us 500.00 MB 500.00 MB 128 + aten::nll_loss_backward 0.69% 1.745ms 2.19% 5.488ms 42.873us 500.00 MB 500.00 MB 128 + aten::_log_softmax_backward_data 14.33% 35.987ms 14.33% 35.987ms 281.146us 500.00 MB 500.00 MB 128 + aten::empty_strided 0.00% 8.894us 0.00% 8.894us 1.779us 32.02 KB 32.02 KB 5 + aten::nll_loss_forward 1.24% 3.107ms 1.24% 3.107ms 24.270us 16.50 KB 16.50 KB 128 + aten::ne 0.01% 21.167us 0.01% 21.167us 21.167us 4.00 KB 4.00 KB 1 + aten::maximum 0.00% 8.209us 0.00% 8.209us 8.209us 8 B 8 B 1 + (198): _run_module_as_main 0.00% 0.500us 100.00% 251.147ms 251.147ms 500.00 MB 0 B 1 + (88): _run_code 0.00% 1.125us 100.00% 251.147ms 251.147ms 500.00 MB 0 B 1 + litgpt/scripts/profile_memory.py(125): 0.00% 0.958us 100.00% 251.145ms 251.145ms 500.00 MB 0 B 1 + litgpt/scripts/profile_memory.py(70): main 0.00% 0.958us 100.00% 251.144ms 251.144ms 500.00 MB 0 B 1 + torch/profiler/profiler.py(1118): __enter__ 0.00% 0.416us 0.05% 133.903us 133.903us 0 B 0 B 1 + torch/profiler/profiler.py(1128): start 0.00% 0.417us 0.05% 133.487us 133.487us 0 B 0 B 1 + torch/profiler/profiler.py(1262): _transit_action 0.00% 7.333us 0.05% 133.070us 133.070us 0 B 0 B 1 +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ +Self CPU time total: 251.147ms + +``` + +### chunk_size=64 + +``` +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + Name Self CPU % Self CPU CPU total % CPU total CPU time avg CPU Mem Self CPU Mem # of Calls +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + aten::cat 8.38% 6.984ms 8.38% 6.991ms 3.496ms 500.02 MB 500.02 MB 2 + aten::_log_softmax 44.11% 36.782ms 44.11% 36.782ms 574.718us 500.00 MB 500.00 MB 64 + aten::nll_loss_backward 1.24% 1.030ms 4.69% 3.908ms 61.061us 500.00 MB 500.00 MB 64 + aten::_log_softmax_backward_data 31.84% 26.553ms 31.84% 26.553ms 414.884us 500.00 MB 500.00 MB 64 + aten::empty_strided 0.00% 2.709us 0.00% 2.709us 0.542us 32.02 KB 32.02 KB 5 + aten::nll_loss_forward 1.30% 1.082ms 1.30% 1.082ms 16.908us 16.25 KB 16.25 KB 64 + aten::ne 0.01% 10.560us 0.01% 10.560us 10.560us 4.00 KB 4.00 KB 1 + aten::maximum 0.00% 2.583us 0.00% 2.583us 2.583us 8 B 8 B 1 + (198): _run_module_as_main 0.00% 0.625us 100.00% 83.393ms 83.393ms 500.00 MB 0 B 1 + (88): _run_code 0.00% 1.334us 100.00% 83.392ms 83.392ms 500.00 MB 0 B 1 + litgpt/scripts/profile_memory.py(125): 0.00% 0.875us 100.00% 83.391ms 83.391ms 500.00 MB 0 B 1 + litgpt/scripts/profile_memory.py(70): main 0.00% 1.166us 100.00% 83.390ms 83.390ms 500.00 MB 0 B 1 + torch/profiler/profiler.py(1118): __enter__ 0.00% 0.292us 0.07% 56.870us 56.870us 0 B 0 B 1 + torch/profiler/profiler.py(1128): start 0.00% 0.292us 0.07% 56.578us 56.578us 0 B 0 B 1 + torch/profiler/profiler.py(1262): _transit_action 0.00% 2.125us 0.07% 56.286us 56.286us 0 B 0 B 1 +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ +Self CPU time total: 83.393ms + +``` + +### chunk_size=128 + +``` +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + Name Self CPU % Self CPU CPU total % CPU total CPU time avg CPU Mem Self CPU Mem # of Calls +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + aten::cat 4.28% 6.221ms 4.28% 6.229ms 3.114ms 500.02 MB 500.02 MB 2 + aten::_log_softmax 51.73% 75.246ms 51.73% 75.246ms 2.351ms 500.00 MB 500.00 MB 32 + aten::nll_loss_backward 0.84% 1.227ms 4.21% 6.122ms 191.301us 500.00 MB 500.00 MB 32 + aten::_log_softmax_backward_data 27.44% 39.912ms 27.44% 39.912ms 1.247ms 500.00 MB 500.00 MB 32 + aten::empty_strided 0.00% 4.916us 0.00% 4.916us 0.983us 32.02 KB 32.02 KB 5 + aten::nll_loss_forward 1.34% 1.950ms 1.34% 1.950ms 60.949us 16.12 KB 16.12 KB 32 + aten::ne 0.01% 10.458us 0.01% 10.458us 10.458us 4.00 KB 4.00 KB 1 + aten::maximum 0.00% 2.833us 0.00% 2.833us 2.833us 8 B 8 B 1 + (198): _run_module_as_main 0.00% 0.584us 100.00% 145.471ms 145.471ms 500.00 MB 0 B 1 + (88): _run_code 0.00% 1.583us 100.00% 145.471ms 145.471ms 500.00 MB 0 B 1 + litgpt/scripts/profile_memory.py(125): 0.00% 1.559us 100.00% 145.469ms 145.469ms 500.00 MB 0 B 1 + litgpt/scripts/profile_memory.py(70): main 0.00% 1.291us 100.00% 145.468ms 145.468ms 500.00 MB 0 B 1 + torch/profiler/profiler.py(1118): __enter__ 0.00% 0.458us 0.04% 58.268us 58.268us 0 B 0 B 1 + torch/profiler/profiler.py(1128): start 0.00% 0.459us 0.04% 57.810us 57.810us 0 B 0 B 1 + torch/profiler/profiler.py(1262): _transit_action 0.00% 1.916us 0.04% 57.351us 57.351us 0 B 0 B 1 +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ +Self CPU time total: 145.471ms + +``` + +### chunk_size=256 + +``` +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + Name Self CPU % Self CPU CPU total % CPU total CPU time avg CPU Mem Self CPU Mem # of Calls +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + aten::cat 6.84% 5.861ms 6.85% 5.869ms 2.934ms 500.02 MB 500.02 MB 2 + aten::_log_softmax 44.01% 37.701ms 44.01% 37.701ms 2.356ms 500.00 MB 500.00 MB 16 + aten::nll_loss_backward 0.59% 504.330us 5.91% 5.065ms 316.553us 500.00 MB 500.00 MB 16 + aten::_log_softmax_backward_data 37.91% 32.479ms 37.91% 32.479ms 2.030ms 500.00 MB 500.00 MB 16 + aten::empty_strided 0.00% 3.209us 0.00% 3.209us 0.642us 32.02 KB 32.02 KB 5 + aten::nll_loss_forward 0.76% 648.975us 0.76% 648.975us 40.561us 16.06 KB 16.06 KB 16 + aten::ne 0.02% 14.060us 0.02% 14.060us 14.060us 4.00 KB 4.00 KB 1 + aten::maximum 0.00% 2.666us 0.00% 2.666us 2.666us 8 B 8 B 1 + (198): _run_module_as_main 0.00% 0.642us 100.00% 85.670ms 85.670ms 500.00 MB 0 B 1 + (88): _run_code 0.00% 1.791us 100.00% 85.669ms 85.669ms 500.00 MB 0 B 1 + litgpt/scripts/profile_memory.py(125): 0.00% 0.917us 100.00% 85.667ms 85.667ms 500.00 MB 0 B 1 + litgpt/scripts/profile_memory.py(70): main 0.00% 1.167us 100.00% 85.666ms 85.666ms 500.00 MB 0 B 1 + torch/profiler/profiler.py(1118): __enter__ 0.00% 0.292us 0.06% 49.893us 49.893us 0 B 0 B 1 + torch/profiler/profiler.py(1128): start 0.00% 0.500us 0.06% 49.601us 49.601us 0 B 0 B 1 + torch/profiler/profiler.py(1262): _transit_action 0.00% 3.125us 0.06% 49.101us 49.101us 0 B 0 B 1 +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ +Self CPU time total: 85.670ms + +``` + +### chunk_size=512 + +``` +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + Name Self CPU % Self CPU CPU total % CPU total CPU time avg CPU Mem Self CPU Mem # of Calls +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + aten::cat 6.20% 6.522ms 6.21% 6.529ms 3.265ms 500.02 MB 500.02 MB 2 + aten::_log_softmax 44.91% 47.232ms 44.91% 47.232ms 5.904ms 500.00 MB 500.00 MB 8 + aten::nll_loss_backward 0.53% 557.234us 8.70% 9.147ms 1.143ms 500.00 MB 500.00 MB 8 + aten::_log_softmax_backward_data 35.12% 36.936ms 35.12% 36.936ms 4.617ms 500.00 MB 500.00 MB 8 + aten::empty_strided 0.00% 3.208us 0.00% 3.208us 0.642us 32.02 KB 32.02 KB 5 + aten::nll_loss_forward 1.25% 1.319ms 1.25% 1.319ms 164.870us 16.03 KB 16.03 KB 8 + aten::ne 0.01% 11.250us 0.01% 11.250us 11.250us 4.00 KB 4.00 KB 1 + aten::maximum 0.00% 2.542us 0.00% 2.542us 2.542us 8 B 8 B 1 + (198): _run_module_as_main 0.00% 0.583us 100.00% 105.179ms 105.179ms 500.00 MB 0 B 1 + (88): _run_code 0.00% 1.417us 100.00% 105.179ms 105.179ms 500.00 MB 0 B 1 + litgpt/scripts/profile_memory.py(125): 0.00% 2.100us 100.00% 105.177ms 105.177ms 500.00 MB 0 B 1 + litgpt/scripts/profile_memory.py(70): main 0.00% 1.209us 100.00% 105.175ms 105.175ms 500.00 MB 0 B 1 + torch/profiler/profiler.py(1118): __enter__ 0.00% 0.375us 0.05% 50.226us 50.226us 0 B 0 B 1 + torch/profiler/profiler.py(1128): start 0.00% 0.292us 0.05% 49.851us 49.851us 0 B 0 B 1 + torch/profiler/profiler.py(1262): _transit_action 0.00% 1.957us 0.05% 49.559us 49.559us 0 B 0 B 1 +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ +Self CPU time total: 105.179ms + +``` diff --git a/docs/profiling/op_table_gpu.md b/docs/profiling/op_table_gpu.md new file mode 100644 index 0000000000..a3accc95f2 --- /dev/null +++ b/docs/profiling/op_table_gpu.md @@ -0,0 +1,140 @@ +### chunk_size=0 + +``` +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + Name Self CPU % Self CPU CPU total % CPU total CPU time avg Self CUDA Self CUDA % CUDA total CUDA time avg CPU Mem Self CPU Mem CUDA Mem Self CUDA Mem # of Calls +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + aten::_log_softmax 19.45% 61.096ms 25.69% 80.681ms 80.681ms 8.403ms 44.19% 25.208ms 25.208ms 0 B 0 B 500.00 MB 500.00 MB 1 + aten::nll_loss_backward 0.84% 2.629ms 1.04% 3.258ms 3.258ms 20.671us 0.11% 2.272ms 2.272ms 0 B 0 B 500.00 MB 500.00 MB 1 + aten::_log_softmax_backward_data 0.03% 82.051us 0.18% 555.226us 555.226us 8.302ms 43.66% 8.302ms 8.302ms 0 B 0 B 500.00 MB 500.00 MB 1 + aten::nll_loss_forward 11.92% 37.441ms 13.72% 43.089ms 43.089ms 36.480us 0.19% 72.960us 72.960us 0 B 0 B 1.00 KB 1.00 KB 1 + aten::empty_strided 0.01% 28.058us 0.01% 28.058us 28.058us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 512 B 512 B 1 +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ +Self CPU time total: 314.105ms +Self CUDA time total: 19.015ms +``` + +### chunk_size=32 + +``` +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + Name Self CPU % Self CPU CPU total % CPU total CPU time avg Self CUDA Self CUDA % CUDA total CUDA time avg CPU Mem Self CPU Mem CUDA Mem Self CUDA Mem # of Calls +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + aten::_log_softmax_backward_data 1.04% 2.074ms 1.54% 3.076ms 24.029us 7.409ms 29.30% 7.409ms 57.881us 0 B 0 B 512.19 MB 512.19 MB 128 + aten::_log_softmax 1.28% 2.548ms 5.72% 11.431ms 89.303us 9.683ms 38.29% 9.751ms 76.182us 0 B 0 B 511.72 MB 511.72 MB 128 + aten::cat 4.51% 9.012ms 8.34% 16.663ms 8.332ms 4.350ms 17.20% 4.353ms 2.176ms 0 B 0 B 500.02 MB 500.02 MB 2 + aten::nll_loss_backward 1.16% 2.316ms 5.77% 11.522ms 90.014us 358.592us 1.42% 3.067ms 23.960us 0 B 0 B 500.00 MB 499.94 MB 128 + aten::nll_loss_forward 1.54% 3.075ms 3.07% 6.136ms 47.941us 444.959us 1.76% 752.096us 5.876us 0 B 0 B 128.00 KB 128.00 KB 128 + aten::empty 0.47% 942.140us 0.47% 942.140us 7.360us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 64.00 KB 64.00 KB 128 + aten::empty_strided 0.02% 45.112us 0.02% 45.112us 15.037us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 33.00 KB 33.00 KB 3 + aten::ne 8.36% 16.711ms 13.75% 27.484ms 27.484ms 3.104us 0.01% 6.208us 6.208us 0 B 0 B 4.00 KB 4.00 KB 1 + aten::div 14.24% 28.450ms 17.45% 34.871ms 17.436ms 5.440us 0.02% 8.224us 4.112us 0 B 0 B 1.00 KB 1.00 KB 2 + aten::maximum 9.12% 18.219ms 15.86% 31.701ms 31.701ms 2.048us 0.01% 4.096us 4.096us 0 B 0 B 512 B 512 B 1 + aten::reshape 0.01% 19.879us 0.03% 53.182us 17.727us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 3 + aten::view 0.02% 33.303us 0.02% 33.303us 11.101us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 3 + aten::split 0.18% 365.100us 0.75% 1.503ms 751.652us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 2 + aten::narrow 0.26% 523.494us 1.29% 2.574ms 6.704us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 384 + aten::slice 0.83% 1.654ms 1.03% 2.051ms 5.340us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 384 +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ +Self CPU time total: 199.854ms +Self CUDA time total: 25.287ms +``` + +### chunk_size=64 + +``` +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + Name Self CPU % Self CPU CPU total % CPU total CPU time avg Self CUDA Self CUDA % CUDA total CUDA time avg CPU Mem Self CPU Mem CUDA Mem Self CUDA Mem # of Calls +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + aten::cat 0.28% 96.535us 2.16% 752.817us 376.408us 4.373ms 16.80% 4.373ms 2.186ms 0 B 0 B 500.02 MB 500.02 MB 2 + aten::_log_softmax 4.39% 1.535ms 38.93% 13.596ms 212.438us 9.258ms 35.57% 9.398ms 146.841us 0 B 0 B 500.00 MB 500.00 MB 64 + aten::_log_softmax_backward_data 1.94% 678.189us 3.11% 1.086ms 16.973us 9.121ms 35.04% 9.121ms 142.523us 0 B 0 B 500.00 MB 500.00 MB 64 + aten::nll_loss_backward 2.38% 832.255us 12.19% 4.258ms 66.530us 201.782us 0.78% 2.852ms 44.561us 0 B 0 B 500.00 MB 499.97 MB 64 + aten::nll_loss_forward 3.20% 1.118ms 7.48% 2.614ms 40.837us 230.776us 0.89% 387.379us 6.053us 0 B 0 B 64.00 KB 64.00 KB 64 + aten::empty_strided 0.06% 21.985us 0.06% 21.985us 7.328us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 33.00 KB 33.00 KB 3 + aten::empty 0.90% 312.827us 0.90% 312.827us 4.888us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 32.00 KB 32.00 KB 64 + aten::ne 0.09% 31.972us 0.12% 40.843us 40.843us 2.752us 0.01% 2.752us 2.752us 0 B 0 B 4.00 KB 4.00 KB 1 + aten::div 0.12% 43.394us 0.19% 64.760us 32.380us 7.103us 0.03% 7.103us 3.551us 0 B 0 B 1.00 KB 1.00 KB 2 + aten::maximum 0.05% 16.679us 0.10% 35.381us 35.381us 2.336us 0.01% 2.336us 2.336us 0 B 0 B 512 B 512 B 1 + aten::reshape 0.06% 21.201us 0.16% 55.516us 18.505us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 3 + aten::view 0.10% 34.315us 0.10% 34.315us 11.438us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 3 + aten::split 0.55% 193.623us 1.95% 679.747us 339.873us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 2 + aten::narrow 0.74% 259.031us 2.29% 799.940us 4.166us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 192 + aten::slice 1.24% 432.266us 1.55% 540.909us 2.817us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 192 +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ +Self CPU time total: 34.928ms +Self CUDA time total: 26.028ms +``` + +### chunk_size=128 + +``` +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + Name Self CPU % Self CPU CPU total % CPU total CPU time avg Self CUDA Self CUDA % CUDA total CUDA time avg CPU Mem Self CPU Mem CUDA Mem Self CUDA Mem # of Calls +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + aten::_log_softmax 1.26% 1.009ms 18.78% 14.995ms 468.591us 10.595ms 37.96% 10.920ms 341.253us 0 B 0 B 512.00 MB 512.00 MB 32 + aten::_log_softmax_backward_data 0.42% 332.375us 0.91% 727.534us 22.735us 9.709ms 34.79% 9.709ms 303.415us 0 B 0 B 512.00 MB 512.00 MB 32 + aten::nll_loss_backward 0.54% 430.166us 2.59% 2.071ms 64.727us 142.686us 0.51% 2.758ms 86.203us 0 B 0 B 512.00 MB 511.98 MB 32 + aten::cat 0.11% 85.666us 0.59% 472.927us 236.463us 4.543ms 16.28% 4.543ms 2.272ms 0 B 0 B 500.02 MB 500.02 MB 2 + aten::empty_strided 0.03% 24.907us 0.03% 24.907us 8.302us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 33.00 KB 33.00 KB 3 + aten::nll_loss_forward 0.86% 684.997us 1.95% 1.554ms 48.563us 155.074us 0.56% 253.569us 7.924us 0 B 0 B 32.00 KB 32.00 KB 32 + aten::empty 0.20% 162.375us 0.20% 162.375us 5.074us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 16.00 KB 16.00 KB 32 + aten::ne 0.04% 29.601us 0.05% 38.377us 38.377us 3.840us 0.01% 3.840us 3.840us 0 B 0 B 4.00 KB 4.00 KB 1 + aten::div 0.05% 40.248us 0.08% 61.577us 30.788us 9.088us 0.03% 9.088us 4.544us 0 B 0 B 1.00 KB 1.00 KB 2 + aten::maximum 0.02% 18.552us 0.03% 26.938us 26.938us 2.912us 0.01% 2.912us 2.912us 0 B 0 B 512 B 512 B 1 +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ +Self CPU time total: 79.861ms +Self CUDA time total: 27.908ms +``` + +### chunk_size=256 + +``` +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + Name Self CPU % Self CPU CPU total % CPU total CPU time avg Self CUDA Self CUDA % CUDA total CUDA time avg CPU Mem Self CPU Mem CUDA Mem Self CUDA Mem # of Calls +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + aten::_log_softmax 1.22% 548.281us 55.12% 24.734ms 1.546ms 9.888ms 38.61% 10.520ms 657.498us 0 B 0 B 512.00 MB 512.00 MB 16 + aten::_log_softmax_backward_data 0.44% 199.317us 1.57% 703.889us 43.993us 8.732ms 34.09% 8.732ms 545.750us 0 B 0 B 512.00 MB 512.00 MB 16 + aten::nll_loss_backward 0.51% 228.647us 2.81% 1.259ms 78.669us 80.159us 0.31% 2.412ms 150.781us 0 B 0 B 512.00 MB 511.98 MB 16 + aten::cat 0.13% 58.337us 0.70% 312.800us 156.400us 4.392ms 17.15% 4.392ms 2.196ms 0 B 0 B 500.02 MB 500.02 MB 2 + aten::empty_strided 0.05% 20.897us 0.05% 20.897us 6.966us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 33.00 KB 33.00 KB 3 + aten::nll_loss_forward 0.75% 336.550us 1.69% 757.204us 47.325us 83.614us 0.33% 135.834us 8.490us 0 B 0 B 24.00 KB 24.00 KB 16 + aten::empty 0.20% 87.874us 0.20% 87.874us 5.492us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 16.00 KB 16.00 KB 16 + aten::ne 0.06% 28.639us 0.08% 37.152us 37.152us 4.192us 0.02% 4.192us 4.192us 0 B 0 B 4.00 KB 4.00 KB 1 + aten::div 0.09% 40.039us 0.13% 59.899us 29.949us 9.632us 0.04% 9.632us 4.816us 0 B 0 B 1.00 KB 1.00 KB 2 + aten::maximum 0.04% 16.688us 0.06% 25.243us 25.243us 3.072us 0.01% 3.072us 3.072us 0 B 0 B 512 B 512 B 1 + aten::reshape 0.04% 19.743us 0.12% 51.890us 17.297us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 3 + aten::view 0.07% 32.147us 0.07% 32.147us 10.716us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 3 + aten::split 0.13% 58.880us 0.43% 191.598us 95.799us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 2 + aten::narrow 0.16% 72.951us 0.49% 219.824us 4.580us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 48 + aten::slice 0.26% 115.594us 0.33% 146.873us 3.060us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 48 +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ +Self CPU time total: 44.868ms +Self CUDA time total: 25.612ms +``` + +### chunk_size=512 + +``` +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + Name Self CPU % Self CPU CPU total % CPU total CPU time avg Self CUDA Self CUDA % CUDA total CUDA time avg CPU Mem Self CPU Mem CUDA Mem Self CUDA Mem # of Calls +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ + aten::cat 0.19% 58.782us 1.95% 600.039us 300.019us 5.011ms 18.47% 5.011ms 2.506ms 0 B 0 B 500.02 MB 500.02 MB 2 + aten::_log_softmax 0.99% 304.731us 19.28% 5.934ms 741.797us 9.958ms 36.71% 11.204ms 1.400ms 0 B 0 B 500.00 MB 500.00 MB 8 + aten::_log_softmax_backward_data 0.36% 109.782us 1.09% 335.018us 41.877us 9.479ms 34.95% 9.479ms 1.185ms 0 B 0 B 500.00 MB 500.00 MB 8 + aten::nll_loss_backward 0.41% 125.890us 2.54% 781.913us 97.739us 47.559us 0.18% 2.553ms 319.171us 0 B 0 B 500.00 MB 499.98 MB 8 + aten::empty_strided 0.07% 20.490us 0.07% 20.490us 6.830us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 33.00 KB 33.00 KB 3 + aten::nll_loss_forward 0.57% 176.829us 1.31% 402.128us 50.266us 46.280us 0.17% 73.260us 9.158us 0 B 0 B 20.00 KB 20.00 KB 8 + aten::empty 0.15% 44.902us 0.15% 44.902us 5.613us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 16.00 KB 16.00 KB 8 + aten::ne 0.10% 30.220us 0.13% 38.826us 38.826us 3.969us 0.01% 3.969us 3.969us 0 B 0 B 4.00 KB 4.00 KB 1 + aten::div 0.13% 40.283us 0.20% 61.670us 30.835us 8.865us 0.03% 8.865us 4.433us 0 B 0 B 1.00 KB 1.00 KB 2 + aten::maximum 0.05% 16.730us 0.08% 24.813us 24.813us 2.977us 0.01% 2.977us 2.977us 0 B 0 B 512 B 512 B 1 + aten::reshape 0.06% 18.841us 0.16% 49.094us 16.365us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 3 + aten::view 0.10% 30.253us 0.10% 30.253us 10.084us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 3 + aten::split 0.11% 34.600us 0.35% 107.196us 53.598us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 2 + aten::narrow 0.13% 39.690us 0.39% 120.553us 5.023us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 24 + aten::slice 0.21% 64.849us 0.26% 80.863us 3.369us 0.000us 0.00% 0.000us 0.000us 0 B 0 B 0 B 0 B 24 +------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ +Self CPU time total: 30.780ms +Self CUDA time total: 27.124ms +``` diff --git a/docs/profiling/peak_memory_vs_chunk_size.png b/docs/profiling/peak_memory_vs_chunk_size.png new file mode 100644 index 0000000000..859f602249 Binary files /dev/null and b/docs/profiling/peak_memory_vs_chunk_size.png differ diff --git a/docs/profiling/peak_memory_vs_chunk_size_gpu.png b/docs/profiling/peak_memory_vs_chunk_size_gpu.png new file mode 100644 index 0000000000..10fc88f013 Binary files /dev/null and b/docs/profiling/peak_memory_vs_chunk_size_gpu.png differ diff --git a/docs/profiling/trace_timeline_chunk128.png b/docs/profiling/trace_timeline_chunk128.png new file mode 100644 index 0000000000..291207ad84 Binary files /dev/null and b/docs/profiling/trace_timeline_chunk128.png differ diff --git a/docs/profiling/trace_timeline_unchunked.png b/docs/profiling/trace_timeline_unchunked.png new file mode 100644 index 0000000000..16c7ce4384 Binary files /dev/null and b/docs/profiling/trace_timeline_unchunked.png differ diff --git a/litgpt/scripts/profile_memory.py b/litgpt/scripts/profile_memory.py new file mode 100644 index 0000000000..825aa3eca1 --- /dev/null +++ b/litgpt/scripts/profile_memory.py @@ -0,0 +1,195 @@ +# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file. + +"""Profiles the memory behavior of `chunked_cross_entropy` (litgpt/utils.py) across a sweep of +`chunk_size` values, to give reproducible evidence for the memory-management concerns raised in +https://github.com/Lightning-AI/litgpt/issues/2190. + +Usage: + python -m litgpt.scripts.profile_memory --output-dir docs/profiling +""" + +import argparse +import json +import tempfile +from pathlib import Path + +import torch + +from litgpt.utils import chunked_cross_entropy + +DEFAULT_CHUNK_SIZES = [0, 32, 64, 128, 256, 512] + + +def _peak_bytes_from_memory_timeline(timeline_path: Path) -> int: + with open(timeline_path) as f: + times, sizes_by_category = json.load(f) + return max((sum(sizes) for sizes in sizes_by_category), default=0) + + +def profile_chunk_size( + chunk_size: int, + batch_size: int, + seq_length: int, + vocab_size: int, + scratch_dir: Path, + device: str = "cpu", + export_memory_plot: Path = None, +): + logits = torch.randn(batch_size, seq_length, vocab_size, requires_grad=True, device=device) + targets = torch.randint(0, vocab_size, (batch_size, seq_length), device=device) + + activities = [torch.profiler.ProfilerActivity.CPU] + if device == "cuda": + activities.append(torch.profiler.ProfilerActivity.CUDA) + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + + with torch.profiler.profile( + activities=activities, + profile_memory=True, + record_shapes=True, + with_stack=export_memory_plot is not None, + ) as prof: + loss = chunked_cross_entropy(logits, targets, chunk_size=chunk_size) + loss.backward() + if device == "cuda": + torch.cuda.synchronize() + + if device == "cuda": + peak_bytes = torch.cuda.max_memory_allocated() + sort_key = "self_cuda_memory_usage" + else: + # the memory timeline JSON is an intermediate artifact used only to compute peak_bytes below; + # it's not committed, so it's written to a scratch dir rather than the (committed) output dir. + timeline_path = scratch_dir / f"memory_timeline_chunk_{chunk_size}.json" + prof.export_memory_timeline(str(timeline_path), device="cpu") + peak_bytes = _peak_bytes_from_memory_timeline(timeline_path) + sort_key = "self_cpu_memory_usage" + + if export_memory_plot is not None: + # .html embeds a base64 PNG of the same memory-timeline plot shown in the PyTorch profiler docs; + # .png is written as its own file for embedding directly in docs/PRs. + html_path = export_memory_plot.with_suffix(".html") + prof.export_memory_timeline(str(html_path), device=f"{device}:0" if device == "cuda" else device) + _png_from_memory_timeline_html(html_path, export_memory_plot) + + table = prof.key_averages().table(sort_by=sort_key, row_limit=15) + return peak_bytes, table + + +def _png_from_memory_timeline_html(html_path: Path, png_path: Path): + import base64 + import re + + html = html_path.read_text() + match = re.search(r"data:image/png;base64,([A-Za-z0-9+/=]+)", html) + png_path.write_bytes(base64.b64decode(match.group(1))) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--chunk-sizes", type=int, nargs="+", default=DEFAULT_CHUNK_SIZES) + parser.add_argument("--batch-size", type=int, default=2) + parser.add_argument("--seq-length", type=int, default=2048) + parser.add_argument("--vocab-size", type=int, default=32000) + parser.add_argument("--output-dir", type=Path, default=Path("docs/profiling")) + parser.add_argument("--device", choices=["cpu", "cuda"], default="cuda" if torch.cuda.is_available() else "cpu") + parser.add_argument( + "--memory-plot-chunk-sizes", + type=int, + nargs="*", + default=[0, 128], + help="chunk sizes to additionally export a memory-timeline plot (PNG) for", + ) + args = parser.parse_args() + + if args.device == "cuda" and not torch.cuda.is_available(): + raise SystemExit("--device cuda requested but no CUDA device is available") + + args.output_dir.mkdir(parents=True, exist_ok=True) + if args.device == "cuda": + print(f"Profiling on {torch.cuda.get_device_name(0)}") + + results = {} + tables = {} + with tempfile.TemporaryDirectory() as scratch_dir: + scratch_dir = Path(scratch_dir) + for chunk_size in args.chunk_sizes: + print( + f"Profiling chunk_size={chunk_size} (B={args.batch_size}, T={args.seq_length}, V={args.vocab_size}, device={args.device})..." + ) + tag = "unchunked" if chunk_size == 0 else f"chunk{chunk_size}" + export_memory_plot = ( + args.output_dir / f"memory_timeline_{tag}.png" if chunk_size in args.memory_plot_chunk_sizes else None + ) + peak_bytes, table = profile_chunk_size( + chunk_size, + args.batch_size, + args.seq_length, + args.vocab_size, + scratch_dir, + device=args.device, + export_memory_plot=export_memory_plot, + ) + results[chunk_size] = peak_bytes + tables[chunk_size] = table + print(f" peak profiler-tracked memory: {peak_bytes / 1e6:.1f} MB") + + (args.output_dir / "op_table.md").write_text( + "\n\n".join(f"### chunk_size={chunk_size}\n\n```\n{table}\n```" for chunk_size, table in tables.items()) + ) + + plot_name = "peak_memory_vs_chunk_size.png" if args.device == "cpu" else "peak_memory_vs_chunk_size_gpu.png" + _plot( + results, + args.output_dir / plot_name, + args.batch_size, + args.seq_length, + args.vocab_size, + device=args.device, + ) + print(f"\nWrote {args.output_dir / plot_name} and {args.output_dir / 'op_table.md'}") + + +def _plot(results: dict, output_path: Path, batch_size: int, seq_length: int, vocab_size: int, device: str = "cpu"): + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + chunk_sizes = list(results.keys()) + peak_mb = [results[c] / 1e6 for c in chunk_sizes] + labels = ["unchunked" if c == 0 else str(c) for c in chunk_sizes] + + surface = "#fcfcfb" + ink_primary = "#0b0b0b" + ink_muted = "#898781" + gridline = "#e1e0d9" + bar_color = "#2a78d6" + + device_label = "CPU" if device == "cpu" else torch.cuda.get_device_name(0) + mem_label = "Peak profiler-tracked CPU memory (MB)" if device == "cpu" else "Peak CUDA memory allocated (MB)" + + fig, ax = plt.subplots(figsize=(7, 4.5)) + fig.patch.set_facecolor(surface) + ax.set_facecolor(surface) + ax.bar(labels, peak_mb, color=bar_color, width=0.6) + ax.set_xlabel("cross_entropy_chunk_size", color=ink_primary) + ax.set_ylabel(mem_label, color=ink_primary) + ax.set_title( + f"chunked_cross_entropy backward-pass memory\n(B={batch_size}, T={seq_length}, V={vocab_size}, {device_label})", + color=ink_primary, + ) + ax.tick_params(colors=ink_muted) + ax.spines["top"].set_visible(False) + ax.spines["right"].set_visible(False) + ax.spines["left"].set_color(gridline) + ax.spines["bottom"].set_color(gridline) + ax.yaxis.grid(True, color=gridline, linewidth=1) + ax.set_axisbelow(True) + fig.tight_layout() + fig.savefig(output_path, dpi=150, facecolor=surface) + + +if __name__ == "__main__": + main()