Skip to content

[DTensor] Add framework adapters for Megatron and vLLM sharding - #1651

Open
vmoens wants to merge 5 commits into
gh/vmoens/91/basefrom
gh/vmoens/91/head
Open

[DTensor] Add framework adapters for Megatron and vLLM sharding#1651
vmoens wants to merge 5 commits into
gh/vmoens/91/basefrom
gh/vmoens/91/head

Conversation

@vmoens

@vmoens vmoens commented Mar 10, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Stack from ghstack (oldest at bottom):

Create tensordict/dtensor_adapters/ package with adapter classes that
construct ShardingDescriptors from framework-specific parameter metadata:

  • MegatronSharding: reads tensor_model_parallel, partition_dim, etc.
  • VLLMSharding: reads output_dim/input_dim for column/row parallel

Each adapter provides describe(name, param) for single parameters and
describe_model(model) for all parameters. Adapters are in separate
submodules to avoid hard dependencies on the respective frameworks.

Made-with: Cursor

@vmoens
Update
80ebceb 
[ghstack-poisoned]
vmoens added a commit that referenced this pull request Mar 10, 2026
@vmoens
[DTensor] Add framework adapters for Megatron and vLLM sharding
5556b38 
Create tensordict/dtensor_adapters/ package with adapter classes that
construct ShardingDescriptors from framework-specific parameter metadata:
- MegatronSharding: reads tensor_model_parallel, partition_dim, etc.
- VLLMSharding: reads output_dim/input_dim for column/row parallel

Each adapter provides describe(name, param) for single parameters and
describe_model(model) for all parameters. Adapters are in separate
submodules to avoid hard dependencies on the respective frameworks.

ghstack-source-id: framework_adapters_placeholder
Made-with: Cursor
ghstack-source-id: cf143d2
Pull-Request: #1651
This was referenced Mar 10, 2026
Open
Open
Open
Open
Open
@github-actions

github-actions Bot commented Mar 10, 2026

Copy link
Copy Markdown
Contributor

PR Title Label Error

Unknown or invalid prefix [DTensor].

Current title: [DTensor] Add framework adapters for Megatron and vLLM sharding

Supported Prefixes

Your PR title must start with exactly one of these prefixes (case-insensitive):

Prefix Label Applied Example
[BugFix] or [Fix] bug [BugFix] Fix memory leak in TensorDict
[Feature] Feature [Feature] Add new storage backend
[Doc] or [Docs] documentation [Doc] Update installation guide
[Refactor] Refactor [Refactor] Clean up module imports
[CI] CI [CI] Fix workflow permissions
[Test] or [Tests] Test [Test] Add unit tests for nn module
[Compile] Compile [Compile] Fix torch.compile issue
[Performance] or [Perf] Performance [Perf] Optimize tensor operations
[Deprecation] Deprecation [Deprecation] Mark old function
[Setup] setup [Setup] Update build configuration
[Distributed] or [Dist] Distributed [Distributed] Add scatter collective
[Benchmark] or [Bench] Benchmarks [Benchmark] Add compile benchmark
[Typing] or [Type] Typing [Typing] Add type stubs
[BC-breaking] or [BC] BC-breaking [BC-breaking] Remove deprecated API
[Formatting] or [Format] Formatting [Format] Fix code style
[Quality] Quality [Quality] Improve error messages

Note: Matching is case-insensitive. Common variations (singular/plural) are supported.

This was referenced Mar 10, 2026
Open
Open
@github-actions github-actions Bot added the Test label Mar 10, 2026
@github-actions

github-actions Bot commented Mar 10, 2026

Copy link
Copy Markdown
Contributor

PR Title Label Error

Unknown or invalid prefix [DTensor].

Current title: [DTensor] Add framework adapters for Megatron and vLLM sharding

Supported Prefixes

Your PR title must start with exactly one of these prefixes (case-insensitive):

Prefix Label Applied Example
[BugFix] or [Fix] bug [BugFix] Fix memory leak in TensorDict
[Feature] Feature [Feature] Add new storage backend
[Doc] or [Docs] documentation [Doc] Update installation guide
[Refactor] Refactor [Refactor] Clean up module imports
[CI] CI [CI] Fix workflow permissions
[Test] or [Tests] Test [Test] Add unit tests for nn module
[Compile] Compile [Compile] Fix torch.compile issue
[Performance] or [Perf] Performance [Perf] Optimize tensor operations
[Deprecation] Deprecation [Deprecation] Mark old function
[Setup] setup [Setup] Update build configuration
[Distributed] or [Dist] Distributed [Distributed] Add scatter collective
[Benchmark] or [Bench] Benchmarks [Benchmark] Add compile benchmark
[Typing] or [Type] Typing [Typing] Add type stubs
[BC-breaking] or [BC] BC-breaking [BC-breaking] Remove deprecated API
[Formatting] or [Format] Formatting [Format] Fix code style
[Quality] Quality [Quality] Improve error messages

Note: Matching is case-insensitive. Common variations (singular/plural) are supported.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Mar 10, 2026
@github-actions

github-actions Bot commented Mar 10, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

$\color{#D29922}\textsf{\Large⚠\kern{0.2cm}\normalsize Warning}$ Result of CPU Benchmark Tests

Total Benchmarks: 261. Improved: $\large\color{#35bf28}26$. Worsened: $\large\color{#d91a1a}5$.

Expand to view detailed results
Name Max Mean Ops Ops on Repo HEAD Change
test_plain_set_nested 40.2620μs 15.0361μs 66.5066 KOps/s 66.9502 KOps/s $\color{#d91a1a}-0.66\%$
test_plain_set_stack_nested 40.9720μs 15.5127μs 64.4633 KOps/s 65.5391 KOps/s $\color{#d91a1a}-1.64\%$
test_plain_set_nested_inplace 78.7040μs 16.8410μs 59.3789 KOps/s 59.3164 KOps/s $\color{#35bf28}+0.11\%$
test_plain_set_stack_nested_inplace 0.1071ms 16.7851μs 59.5768 KOps/s 59.5045 KOps/s $\color{#35bf28}+0.12\%$
test_items 31.4220μs 6.0593μs 165.0361 KOps/s 165.6368 KOps/s $\color{#d91a1a}-0.36\%$
test_items_nested 0.5251ms 0.4647ms 2.1519 KOps/s 2.1196 KOps/s $\color{#35bf28}+1.52\%$
test_items_nested_locked 0.5201ms 0.4675ms 2.1392 KOps/s 2.1231 KOps/s $\color{#35bf28}+0.76\%$
test_items_nested_leaf 0.1650ms 97.7323μs 10.2320 KOps/s 10.1318 KOps/s $\color{#35bf28}+0.99\%$
test_items_stack_nested 0.5879ms 0.4634ms 2.1580 KOps/s 2.1167 KOps/s $\color{#35bf28}+1.95\%$
test_items_stack_nested_leaf 0.1226ms 96.9582μs 10.3137 KOps/s 10.2753 KOps/s $\color{#35bf28}+0.37\%$
test_items_stack_nested_locked 0.7178ms 0.4593ms 2.1774 KOps/s 2.1485 KOps/s $\color{#35bf28}+1.35\%$
test_keys 33.4620μs 4.2152μs 237.2355 KOps/s 220.5704 KOps/s $\textbf{\color{#35bf28}+7.56\%}$
test_keys_nested 0.2065ms 0.1298ms 7.7019 KOps/s 7.6660 KOps/s $\color{#35bf28}+0.47\%$
test_keys_nested_locked 0.7806ms 0.1390ms 7.1928 KOps/s 7.2835 KOps/s $\color{#d91a1a}-1.24\%$
test_keys_nested_leaf 0.1471ms 0.1210ms 8.2659 KOps/s 8.3452 KOps/s $\color{#d91a1a}-0.95\%$
test_keys_stack_nested 0.1623ms 0.1306ms 7.6582 KOps/s 7.7338 KOps/s $\color{#d91a1a}-0.98\%$
test_keys_stack_nested_leaf 0.2125ms 0.1207ms 8.2852 KOps/s 8.3853 KOps/s $\color{#d91a1a}-1.19\%$
test_keys_stack_nested_locked 0.2001ms 0.1376ms 7.2653 KOps/s 7.1860 KOps/s $\color{#35bf28}+1.10\%$
test_values 0.1093ms 1.0143μs 985.9349 KOps/s 977.0376 KOps/s $\color{#35bf28}+0.91\%$
test_values_nested 74.7740μs 51.9471μs 19.2504 KOps/s 19.3063 KOps/s $\color{#d91a1a}-0.29\%$
test_values_nested_locked 80.7230μs 55.8354μs 17.9098 KOps/s 17.7627 KOps/s $\color{#35bf28}+0.83\%$
test_values_nested_leaf 0.1431ms 59.4559μs 16.8192 KOps/s 16.5813 KOps/s $\color{#35bf28}+1.43\%$
test_values_stack_nested 90.8340μs 52.2100μs 19.1534 KOps/s 18.9130 KOps/s $\color{#35bf28}+1.27\%$
test_values_stack_nested_leaf 89.5350μs 59.7596μs 16.7337 KOps/s 16.5761 KOps/s $\color{#35bf28}+0.95\%$
test_values_stack_nested_locked 93.0140μs 55.5076μs 18.0155 KOps/s 17.8637 KOps/s $\color{#35bf28}+0.85\%$
test_membership 5.4368μs 0.8272μs 1.2089 MOps/s 1.1644 MOps/s $\color{#35bf28}+3.83\%$
test_membership_nested 25.5820μs 2.8410μs 351.9849 KOps/s 343.6517 KOps/s $\color{#35bf28}+2.42\%$
test_membership_nested_leaf 32.3310μs 2.8839μs 346.7478 KOps/s 358.0727 KOps/s $\color{#d91a1a}-3.16\%$
test_membership_stacked_nested 23.3710μs 2.8887μs 346.1744 KOps/s 348.7288 KOps/s $\color{#d91a1a}-0.73\%$
test_membership_stacked_nested_leaf 27.8210μs 2.8879μs 346.2719 KOps/s 351.0304 KOps/s $\color{#d91a1a}-1.36\%$
test_membership_nested_last 23.3710μs 4.3206μs 231.4518 KOps/s 231.2013 KOps/s $\color{#35bf28}+0.11\%$
test_membership_nested_leaf_last 34.8210μs 4.3257μs 231.1773 KOps/s 227.2521 KOps/s $\color{#35bf28}+1.73\%$
test_membership_stacked_nested_last 23.2510μs 4.3190μs 231.5371 KOps/s 226.2918 KOps/s $\color{#35bf28}+2.32\%$
test_membership_stacked_nested_leaf_last 30.9810μs 4.3641μs 229.1440 KOps/s 227.9179 KOps/s $\color{#35bf28}+0.54\%$
test_nested_getleaf 46.8820μs 21.4089μs 46.7094 KOps/s 46.3037 KOps/s $\color{#35bf28}+0.88\%$
test_nested_get 52.0520μs 20.2462μs 49.3920 KOps/s 48.2703 KOps/s $\color{#35bf28}+2.32\%$
test_stacked_getleaf 0.1073ms 21.3562μs 46.8248 KOps/s 46.4472 KOps/s $\color{#35bf28}+0.81\%$
test_stacked_get 51.0120μs 20.4415μs 48.9202 KOps/s 48.4207 KOps/s $\color{#35bf28}+1.03\%$
test_nested_getitemleaf 59.6930μs 21.8534μs 45.7594 KOps/s 45.2126 KOps/s $\color{#35bf28}+1.21\%$
test_nested_getitem 47.2020μs 20.8038μs 48.0682 KOps/s 47.6766 KOps/s $\color{#35bf28}+0.82\%$
test_stacked_getitemleaf 52.3520μs 21.8366μs 45.7947 KOps/s 45.5616 KOps/s $\color{#35bf28}+0.51\%$
test_stacked_getitem 46.0520μs 20.5237μs 48.7241 KOps/s 47.6795 KOps/s $\color{#35bf28}+2.19\%$
test_lock_nested 8.2417ms 0.4886ms 2.0466 KOps/s 2.1150 KOps/s $\color{#d91a1a}-3.24\%$
test_lock_stack_nested 0.5627ms 0.4813ms 2.0779 KOps/s 2.0718 KOps/s $\color{#35bf28}+0.29\%$
test_unlock_nested 0.4909ms 0.3887ms 2.5728 KOps/s 2.5777 KOps/s $\color{#d91a1a}-0.19\%$
test_unlock_stack_nested 0.4617ms 0.3872ms 2.5829 KOps/s 2.5526 KOps/s $\color{#35bf28}+1.19\%$
test_flatten_speed 0.1754ms 0.1229ms 8.1356 KOps/s 8.1486 KOps/s $\color{#d91a1a}-0.16\%$
test_unflatten_speed 0.6082ms 0.5744ms 1.7411 KOps/s 1.7124 KOps/s $\color{#35bf28}+1.68\%$
test_common_ops 0.8682ms 0.6988ms 1.4311 KOps/s 1.4233 KOps/s $\color{#35bf28}+0.55\%$
test_creation 62.6530μs 3.1112μs 321.4216 KOps/s 319.5978 KOps/s $\color{#35bf28}+0.57\%$
test_creation_empty 65.2530μs 7.0459μs 141.9266 KOps/s 142.0766 KOps/s $\color{#d91a1a}-0.11\%$
test_creation_nested_1 41.4220μs 11.6561μs 85.7917 KOps/s 85.4883 KOps/s $\color{#35bf28}+0.35\%$
test_creation_nested_2 54.2920μs 13.3312μs 75.0122 KOps/s 74.6813 KOps/s $\color{#35bf28}+0.44\%$
test_creation_many_keys[10] 46.7420μs 21.0654μs 47.4711 KOps/s 47.0852 KOps/s $\color{#35bf28}+0.82\%$
test_creation_many_keys[50] 0.1207ms 90.8976μs 11.0014 KOps/s 11.0873 KOps/s $\color{#d91a1a}-0.77\%$
test_creation_many_keys[100] 0.2022ms 0.1782ms 5.6101 KOps/s 5.5898 KOps/s $\color{#35bf28}+0.36\%$
test_creation_nested_many_keys[10] 72.5530μs 45.3796μs 22.0364 KOps/s 22.0893 KOps/s $\color{#d91a1a}-0.24\%$
test_creation_nested_many_keys[50] 0.2648ms 0.1856ms 5.3873 KOps/s 5.4402 KOps/s $\color{#d91a1a}-0.97\%$
test_clone 39.4420μs 13.3002μs 75.1868 KOps/s 73.7424 KOps/s $\color{#35bf28}+1.96\%$
test_getitem[int] 1.5608ms 15.0841μs 66.2950 KOps/s 60.3660 KOps/s $\textbf{\color{#35bf28}+9.82\%}$
test_getitem[slice_int] 0.1406ms 24.0870μs 41.5162 KOps/s 42.8147 KOps/s $\color{#d91a1a}-3.03\%$
test_getitem[range] 0.2057ms 63.3687μs 15.7807 KOps/s 15.7891 KOps/s $\color{#d91a1a}-0.05\%$
test_getitem[tuple] 0.1465ms 24.3762μs 41.0236 KOps/s 42.0483 KOps/s $\color{#d91a1a}-2.44\%$
test_getitem[list] 0.1916ms 58.0775μs 17.2184 KOps/s 17.2408 KOps/s $\color{#d91a1a}-0.13\%$
test_setitem_dim[int] 56.2720μs 26.0459μs 38.3938 KOps/s 37.9544 KOps/s $\color{#35bf28}+1.16\%$
test_setitem_dim[slice_int] 66.8830μs 43.3553μs 23.0652 KOps/s 22.7947 KOps/s $\color{#35bf28}+1.19\%$
test_setitem_dim[range] 0.1250ms 94.7451μs 10.5546 KOps/s 10.5974 KOps/s $\color{#d91a1a}-0.40\%$
test_setitem_dim[tuple] 65.8330μs 39.4798μs 25.3294 KOps/s 24.9669 KOps/s $\color{#35bf28}+1.45\%$
test_setitem 51.3320μs 17.4064μs 57.4502 KOps/s 54.5323 KOps/s $\textbf{\color{#35bf28}+5.35\%}$
test_set 49.1520μs 16.9964μs 58.8359 KOps/s 56.6823 KOps/s $\color{#35bf28}+3.80\%$
test_set_shared 0.5714ms 0.2041ms 4.9006 KOps/s 4.8672 KOps/s $\color{#35bf28}+0.69\%$
test_update 0.1920ms 21.6301μs 46.2320 KOps/s 45.7773 KOps/s $\color{#35bf28}+0.99\%$
test_update_nested 71.9730μs 34.2823μs 29.1696 KOps/s 28.9346 KOps/s $\color{#35bf28}+0.81\%$
test_update__nested 0.4701ms 33.8256μs 29.5634 KOps/s 28.7072 KOps/s $\color{#35bf28}+2.98\%$
test_set_nested 49.2420μs 19.0003μs 52.6307 KOps/s 52.2819 KOps/s $\color{#35bf28}+0.67\%$
test_set_nested_new 57.9630μs 23.2132μs 43.0790 KOps/s 41.4349 KOps/s $\color{#35bf28}+3.97\%$
test_select 73.8940μs 41.3290μs 24.1961 KOps/s 24.8186 KOps/s $\color{#d91a1a}-2.51\%$
test_select_nested 0.1068ms 73.2937μs 13.6437 KOps/s 13.7162 KOps/s $\color{#d91a1a}-0.53\%$
test_exclude_nested 0.1580ms 90.7639μs 11.0176 KOps/s 10.7872 KOps/s $\color{#35bf28}+2.14\%$
test_empty[True] 0.4813ms 0.3960ms 2.5251 KOps/s 2.4857 KOps/s $\color{#35bf28}+1.58\%$
test_empty[False] 10.3080μs 1.2874μs 776.7310 KOps/s 756.7954 KOps/s $\color{#35bf28}+2.63\%$
test_to 0.1084ms 71.8995μs 13.9083 KOps/s 13.8847 KOps/s $\color{#35bf28}+0.17\%$
test_to_nonblocking 0.1080ms 66.7650μs 14.9779 KOps/s 14.8002 KOps/s $\color{#35bf28}+1.20\%$
test_unbind_speed 0.3710ms 0.3345ms 2.9895 KOps/s 3.0016 KOps/s $\color{#d91a1a}-0.40\%$
test_unbind_speed_stack0 0.3988ms 0.3293ms 3.0363 KOps/s 3.0218 KOps/s $\color{#35bf28}+0.48\%$
test_unbind_speed_stack1 0.1050s 0.8316ms 1.2025 KOps/s 1.1863 KOps/s $\color{#35bf28}+1.37\%$
test_split 0.1051s 1.2613ms 792.8349 Ops/s 786.8720 Ops/s $\color{#35bf28}+0.76\%$
test_chunk 0.1045s 1.2094ms 826.8601 Ops/s 921.4780 Ops/s $\textbf{\color{#d91a1a}-10.27\%}$
test_to_cpu_blocking 19.8707ms 19.5808ms 51.0704 Ops/s 34.8827 Ops/s $\textbf{\color{#35bf28}+46.41\%}$
test_to_cpu_global_sync 11.5498ms 11.2450ms 88.9287 Ops/s 79.3449 Ops/s $\textbf{\color{#35bf28}+12.08\%}$
test_to_cpu_event_sync 12.4998ms 12.2354ms 81.7298 Ops/s 81.0790 Ops/s $\color{#35bf28}+0.80\%$
test_to_cpu_default 0.1167s 13.5257ms 73.9331 Ops/s 80.7277 Ops/s $\textbf{\color{#d91a1a}-8.42\%}$
test_consolidate[False-None] 4.2575ms 4.1724ms 239.6694 Ops/s 217.5824 Ops/s $\textbf{\color{#35bf28}+10.15\%}$
test_consolidate[default-None] 2.1079ms 1.9893ms 502.6905 Ops/s 489.7343 Ops/s $\color{#35bf28}+2.65\%$
test_consolidate[reduce-overhead-None] 2.0394ms 1.9301ms 518.0975 Ops/s 509.3891 Ops/s $\color{#35bf28}+1.71\%$
test_consolidate_njt[False-None] 8.9905ms 8.5387ms 117.1136 Ops/s 117.1665 Ops/s $\color{#d91a1a}-0.05\%$
test_to[False-False-None] 2.2339ms 2.0777ms 481.3113 Ops/s 477.9272 Ops/s $\color{#35bf28}+0.71\%$
test_to[True-False-None] 2.1732ms 1.9373ms 516.1887 Ops/s 513.7259 Ops/s $\color{#35bf28}+0.48\%$
test_to[within-False-None] 6.4140ms 6.1653ms 162.1994 Ops/s 161.8228 Ops/s $\color{#35bf28}+0.23\%$
test_to[True-default-None] 9.3500ms 8.8835ms 112.5678 Ops/s 110.9457 Ops/s $\color{#35bf28}+1.46\%$
test_to_njt[False-False-None] 8.5149ms 8.3882ms 119.2151 Ops/s 117.1962 Ops/s $\color{#35bf28}+1.72\%$
test_to_njt[True-False-None] 7.4027ms 7.0385ms 142.0756 Ops/s 143.9326 Ops/s $\color{#d91a1a}-1.29\%$
test_to_njt[within-False-None] 16.3175ms 15.4950ms 64.5369 Ops/s 63.9692 Ops/s $\color{#35bf28}+0.89\%$
test_creation[device0] 0.4526ms 0.1154ms 8.6687 KOps/s 8.7129 KOps/s $\color{#d91a1a}-0.51\%$
test_creation_from_tensor 0.3965ms 0.1165ms 8.5820 KOps/s 8.7401 KOps/s $\color{#d91a1a}-1.81\%$
test_add_one[memmap_tensor0] 0.3299ms 6.5585μs 152.4745 KOps/s 152.8875 KOps/s $\color{#d91a1a}-0.27\%$
test_contiguous[memmap_tensor0] 27.2010μs 0.6715μs 1.4891 MOps/s 2.0066 MOps/s $\textbf{\color{#d91a1a}-25.79\%}$
test_stack[memmap_tensor0] 23.7510μs 4.6010μs 217.3459 KOps/s 220.3181 KOps/s $\color{#d91a1a}-1.35\%$
test_memmaptd_index 1.2688ms 0.2713ms 3.6860 KOps/s 3.7229 KOps/s $\color{#d91a1a}-0.99\%$
test_memmaptd_index_astensor 0.5421ms 0.3784ms 2.6430 KOps/s 2.6998 KOps/s $\color{#d91a1a}-2.10\%$
test_memmaptd_index_op 0.7964ms 0.6185ms 1.6167 KOps/s 1.6036 KOps/s $\color{#35bf28}+0.82\%$
test_serialize_model 0.1379s 0.1363s 7.3364 Ops/s 7.3715 Ops/s $\color{#d91a1a}-0.48\%$
test_serialize_model_pickle 1.3488s 1.2106s 0.8261 Ops/s 0.8256 Ops/s $\color{#35bf28}+0.05\%$
test_serialize_weights 0.1402s 0.1352s 7.3945 Ops/s 7.4064 Ops/s $\color{#d91a1a}-0.16\%$
test_serialize_weights_returnearly 0.4297s 86.9674ms 11.4986 Ops/s 11.8527 Ops/s $\color{#d91a1a}-2.99\%$
test_serialize_weights_pickle 1.3531s 1.2199s 0.8197 Ops/s 0.8245 Ops/s $\color{#d91a1a}-0.58\%$
test_reshape_pytree 0.2039ms 32.4278μs 30.8377 KOps/s 30.4125 KOps/s $\color{#35bf28}+1.40\%$
test_reshape_td 81.1140μs 45.8921μs 21.7902 KOps/s 22.3341 KOps/s $\color{#d91a1a}-2.43\%$
test_view_pytree 0.2202ms 31.8587μs 31.3886 KOps/s 30.5964 KOps/s $\color{#35bf28}+2.59\%$
test_view_td 0.1069ms 54.4964μs 18.3498 KOps/s 18.8777 KOps/s $\color{#d91a1a}-2.80\%$
test_unbind_pytree 0.2299ms 36.3111μs 27.5398 KOps/s 27.3644 KOps/s $\color{#35bf28}+0.64\%$
test_unbind_td 0.1511ms 49.9148μs 20.0341 KOps/s 20.1551 KOps/s $\color{#d91a1a}-0.60\%$
test_split_pytree 0.1984ms 42.0615μs 23.7747 KOps/s 23.5262 KOps/s $\color{#35bf28}+1.06\%$
test_split_td 0.1473ms 64.7833μs 15.4361 KOps/s 15.4380 KOps/s $\color{#d91a1a}-0.01\%$
test_add_pytree 0.2348ms 41.5207μs 24.0844 KOps/s 23.3133 KOps/s $\color{#35bf28}+3.31\%$
test_add_td 0.2543ms 55.8398μs 17.9084 KOps/s 17.3873 KOps/s $\color{#35bf28}+3.00\%$
test_compile_add_one_nested[tensordict-compile] 0.1904ms 0.1408ms 7.1002 KOps/s 6.8605 KOps/s $\color{#35bf28}+3.49\%$
test_compile_add_one_nested[tensordict-eager] 0.4813ms 0.2020ms 4.9512 KOps/s 5.0190 KOps/s $\color{#d91a1a}-1.35\%$
test_compile_add_one_nested[pytree-compile] 0.2802ms 0.1131ms 8.8432 KOps/s 9.1696 KOps/s $\color{#d91a1a}-3.56\%$
test_compile_add_one_nested[pytree-eager] 0.4445ms 0.1778ms 5.6254 KOps/s 5.5736 KOps/s $\color{#35bf28}+0.93\%$
test_compile_copy_nested[tensordict-compile] 0.3091ms 9.9131μs 100.8764 KOps/s 92.0889 KOps/s $\textbf{\color{#35bf28}+9.54\%}$
test_compile_copy_nested[tensordict-eager] 82.6940μs 54.2435μs 18.4354 KOps/s 18.3642 KOps/s $\color{#35bf28}+0.39\%$
test_compile_copy_nested[pytree-compile] 39.9920μs 9.7501μs 102.5629 KOps/s 102.1191 KOps/s $\color{#35bf28}+0.43\%$
test_compile_copy_nested[pytree-eager] 0.4711ms 67.9872μs 14.7087 KOps/s 14.2709 KOps/s $\color{#35bf28}+3.07\%$
test_compile_add_one_flat[tensordict-compile] 0.2140ms 0.1753ms 5.7043 KOps/s 5.2825 KOps/s $\textbf{\color{#35bf28}+7.98\%}$
test_compile_add_one_flat[tensordict-eager] 0.3449ms 0.2834ms 3.5290 KOps/s 3.5202 KOps/s $\color{#35bf28}+0.25\%$
test_compile_add_one_flat[tensorclass-compile] 0.2459ms 0.1171ms 8.5391 KOps/s 8.0409 KOps/s $\textbf{\color{#35bf28}+6.20\%}$
test_compile_add_one_flat[tensorclass-eager] 0.1105ms 73.2927μs 13.6439 KOps/s 13.4919 KOps/s $\color{#35bf28}+1.13\%$
test_compile_add_one_flat[pytree-compile] 0.2236ms 0.1576ms 6.3461 KOps/s 5.9257 KOps/s $\textbf{\color{#35bf28}+7.09\%}$
test_compile_add_one_flat[pytree-eager] 0.8015ms 0.5190ms 1.9268 KOps/s 1.9061 KOps/s $\color{#35bf28}+1.08\%$
test_compile_add_self_flat[tensordict-eager] 0.4725ms 0.3326ms 3.0065 KOps/s 2.9518 KOps/s $\color{#35bf28}+1.85\%$
test_compile_add_self_flat[tensordict-compile] 0.2289ms 0.1790ms 5.5872 KOps/s 5.2589 KOps/s $\textbf{\color{#35bf28}+6.24\%}$
test_compile_add_self_flat[tensorclass-eager] 0.1599ms 88.4824μs 11.3017 KOps/s 11.0904 KOps/s $\color{#35bf28}+1.91\%$
test_compile_add_self_flat[tensorclass-compile] 0.1848ms 0.1202ms 8.3198 KOps/s 8.0164 KOps/s $\color{#35bf28}+3.78\%$
test_compile_add_self_flat[pytree-eager] 0.7097ms 0.4307ms 2.3219 KOps/s 2.2939 KOps/s $\color{#35bf28}+1.22\%$
test_compile_add_self_flat[pytree-compile] 0.2030ms 0.1592ms 6.2821 KOps/s 6.1561 KOps/s $\color{#35bf28}+2.05\%$
test_compile_copy_flat[tensordict-compile] 37.2320μs 13.2406μs 75.5255 KOps/s 73.5831 KOps/s $\color{#35bf28}+2.64\%$
test_compile_copy_flat[tensordict-eager] 91.3350μs 40.5028μs 24.6896 KOps/s 23.9362 KOps/s $\color{#35bf28}+3.15\%$
test_compile_copy_flat[pytree-compile] 50.5930μs 10.7394μs 93.1148 KOps/s 92.9408 KOps/s $\color{#35bf28}+0.19\%$
test_compile_copy_flat[pytree-eager] 0.3945ms 52.7062μs 18.9731 KOps/s 19.1700 KOps/s $\color{#d91a1a}-1.03\%$
test_compile_assign_and_add[tensordict-compile] 1.9845ms 0.1731ms 5.7754 KOps/s 5.5694 KOps/s $\color{#35bf28}+3.70\%$
test_compile_assign_and_add[tensordict-eager] 3.7596ms 3.3040ms 302.6667 Ops/s 301.8986 Ops/s $\color{#35bf28}+0.25\%$
test_compile_assign_and_add[pytree-compile] 1.9781ms 0.1608ms 6.2208 KOps/s 5.9827 KOps/s $\color{#35bf28}+3.98\%$
test_compile_assign_and_add[pytree-eager] 3.1931ms 2.7712ms 360.8554 Ops/s 357.7852 Ops/s $\color{#35bf28}+0.86\%$
test_compile_indexing[tensor-tensordict-compile] 0.1541ms 0.1088ms 9.1951 KOps/s 8.9265 KOps/s $\color{#35bf28}+3.01\%$
test_compile_indexing[tensor-tensordict-eager] 0.5245ms 73.2790μs 13.6465 KOps/s 12.9478 KOps/s $\textbf{\color{#35bf28}+5.40\%}$
test_compile_indexing[tensor-tensorclass-compile] 0.5924ms 95.7214μs 10.4470 KOps/s 10.1873 KOps/s $\color{#35bf28}+2.55\%$
test_compile_indexing[tensor-tensorclass-eager] 0.4741ms 43.6753μs 22.8963 KOps/s 21.0178 KOps/s $\textbf{\color{#35bf28}+8.94\%}$
test_compile_indexing[tensor-pytree-compile] 0.5574ms 97.0947μs 10.2992 KOps/s 10.1380 KOps/s $\color{#35bf28}+1.59\%$
test_compile_indexing[tensor-pytree-eager] 0.4816ms 43.7690μs 22.8472 KOps/s 22.4369 KOps/s $\color{#35bf28}+1.83\%$
test_compile_indexing[slice-tensordict-compile] 0.5021ms 56.6847μs 17.6414 KOps/s 17.5493 KOps/s $\color{#35bf28}+0.52\%$
test_compile_indexing[slice-tensordict-eager] 0.2287ms 27.1472μs 36.8363 KOps/s 35.6821 KOps/s $\color{#35bf28}+3.23\%$
test_compile_indexing[slice-tensorclass-compile] 0.5267ms 45.0036μs 22.2205 KOps/s 22.2436 KOps/s $\color{#d91a1a}-0.10\%$
test_compile_indexing[slice-tensorclass-eager] 0.4561ms 22.3317μs 44.7794 KOps/s 43.6253 KOps/s $\color{#35bf28}+2.65\%$
test_compile_indexing[slice-pytree-compile] 0.4865ms 44.5675μs 22.4379 KOps/s 21.7164 KOps/s $\color{#35bf28}+3.32\%$
test_compile_indexing[slice-pytree-eager] 0.4443ms 22.3394μs 44.7640 KOps/s 43.8818 KOps/s $\color{#35bf28}+2.01\%$
test_compile_indexing[int-tensordict-compile] 0.4883ms 57.1259μs 17.5052 KOps/s 17.1283 KOps/s $\color{#35bf28}+2.20\%$
test_compile_indexing[int-tensordict-eager] 0.4554ms 27.3462μs 36.5681 KOps/s 36.0436 KOps/s $\color{#35bf28}+1.46\%$
test_compile_indexing[int-tensorclass-compile] 81.3840μs 45.0828μs 22.1814 KOps/s 21.6886 KOps/s $\color{#35bf28}+2.27\%$
test_compile_indexing[int-tensorclass-eager] 0.4690ms 22.1331μs 45.1812 KOps/s 43.6301 KOps/s $\color{#35bf28}+3.56\%$
test_compile_indexing[int-pytree-compile] 0.4789ms 45.1782μs 22.1346 KOps/s 21.6941 KOps/s $\color{#35bf28}+2.03\%$
test_compile_indexing[int-pytree-eager] 0.4421ms 22.2847μs 44.8739 KOps/s 43.6745 KOps/s $\color{#35bf28}+2.75\%$
test_compile_replace[single-eager] 0.4961ms 47.7920μs 20.9240 KOps/s 20.0228 KOps/s $\color{#35bf28}+4.50\%$
test_compile_replace[single-compile] 0.1672ms 0.1050ms 9.5210 KOps/s 9.5387 KOps/s $\color{#d91a1a}-0.18\%$
test_compile_replace[multi-eager] 1.0285ms 0.5810ms 1.7213 KOps/s 1.7463 KOps/s $\color{#d91a1a}-1.43\%$
test_compile_replace[multi-compile] 0.1788ms 0.1108ms 9.0276 KOps/s 8.6202 KOps/s $\color{#35bf28}+4.73\%$
test_compile_tc_getattr_20[eager] 0.6019ms 0.1653ms 6.0498 KOps/s 6.1215 KOps/s $\color{#d91a1a}-1.17\%$
test_compile_tc_getattr_20[compile] 0.5539ms 0.1185ms 8.4415 KOps/s 7.9614 KOps/s $\textbf{\color{#35bf28}+6.03\%}$
test_compile_clone_shallow[20-eager] 0.4414ms 19.2050μs 52.0698 KOps/s 51.3964 KOps/s $\color{#35bf28}+1.31\%$
test_compile_clone_shallow[20-compile] 0.4483ms 11.2239μs 89.0956 KOps/s 85.0216 KOps/s $\color{#35bf28}+4.79\%$
test_compile_clone_shallow[40-eager] 72.4430μs 33.9537μs 29.4519 KOps/s 29.6070 KOps/s $\color{#d91a1a}-0.52\%$
test_compile_clone_shallow[40-compile] 37.8020μs 12.2196μs 81.8357 KOps/s 67.6077 KOps/s $\textbf{\color{#35bf28}+21.05\%}$
test_compile_clone_shallow[80-eager] 0.4933ms 63.3159μs 15.7938 KOps/s 15.8148 KOps/s $\color{#d91a1a}-0.13\%$
test_compile_clone_shallow[80-compile] 0.4439ms 14.7487μs 67.8025 KOps/s 65.1705 KOps/s $\color{#35bf28}+4.04\%$
test_compile_update_inplace[eager] 0.4978ms 58.6703μs 17.0444 KOps/s 16.8459 KOps/s $\color{#35bf28}+1.18\%$
test_compile_update_inplace[compile] 0.5711ms 0.1384ms 7.2236 KOps/s 6.9224 KOps/s $\color{#35bf28}+4.35\%$
test_mod_add[eager] 0.4931ms 48.3305μs 20.6909 KOps/s 20.2944 KOps/s $\color{#35bf28}+1.95\%$
test_mod_add[compile] 0.5953ms 0.1034ms 9.6749 KOps/s 9.3815 KOps/s $\color{#35bf28}+3.13\%$
test_mod_add[compile-overhead] 0.2336ms 0.1465ms 6.8246 KOps/s 6.6217 KOps/s $\color{#35bf28}+3.06\%$
test_mod_wrap[eager] 0.7405ms 0.2884ms 3.4675 KOps/s 3.4253 KOps/s $\color{#35bf28}+1.23\%$
test_mod_wrap[compile] 0.4152ms 0.3450ms 2.8988 KOps/s 2.7928 KOps/s $\color{#35bf28}+3.80\%$
test_mod_wrap[compile-overhead] 7.2738ms 4.0188ms 248.8315 Ops/s 254.0651 Ops/s $\color{#d91a1a}-2.06\%$
test_mod_wrap_and_backward[eager] 1.6078ms 1.4949ms 668.9374 Ops/s 661.7822 Ops/s $\color{#35bf28}+1.08\%$
test_mod_wrap_and_backward[compile] 1.7055ms 1.5459ms 646.8868 Ops/s 689.9235 Ops/s $\textbf{\color{#d91a1a}-6.24\%}$
test_mod_wrap_and_backward[compile-overhead] 1.4508ms 0.9872ms 1.0130 KOps/s 1.1037 KOps/s $\textbf{\color{#d91a1a}-8.21\%}$
test_seq_add[eager] 0.5913ms 0.1522ms 6.5684 KOps/s 6.4124 KOps/s $\color{#35bf28}+2.43\%$
test_seq_add[compile] 0.6003ms 0.1112ms 8.9955 KOps/s 8.3152 KOps/s $\textbf{\color{#35bf28}+8.18\%}$
test_seq_add[compile-overhead] 0.6000ms 0.1519ms 6.5832 KOps/s 6.2364 KOps/s $\textbf{\color{#35bf28}+5.56\%}$
test_seq_wrap[eager] 0.9720ms 0.5227ms 1.9133 KOps/s 1.8482 KOps/s $\color{#35bf28}+3.52\%$
test_seq_wrap[compile] 0.8503ms 0.3627ms 2.7574 KOps/s 2.6627 KOps/s $\color{#35bf28}+3.56\%$
test_seq_wrap[compile-overhead] 0.3250ms 0.2633ms 3.7973 KOps/s 3.6790 KOps/s $\color{#35bf28}+3.22\%$
test_func_call_runtime[False-eager] 1.3006ms 0.8499ms 1.1766 KOps/s 1.1322 KOps/s $\color{#35bf28}+3.92\%$
test_func_call_runtime[False-compile] 1.3952ms 0.8977ms 1.1139 KOps/s 1.0875 KOps/s $\color{#35bf28}+2.43\%$
test_func_call_runtime[False-compile-overhead] 0.9064ms 0.4611ms 2.1687 KOps/s 2.1415 KOps/s $\color{#35bf28}+1.27\%$
test_func_call_runtime[True-eager] 1.5240ms 1.0803ms 925.7053 Ops/s 920.4890 Ops/s $\color{#35bf28}+0.57\%$
test_func_call_runtime[True-compile] 1.3669ms 0.9117ms 1.0969 KOps/s 1.0808 KOps/s $\color{#35bf28}+1.49\%$
test_func_call_runtime[True-compile-overhead] 0.5159ms 0.4722ms 2.1179 KOps/s 2.0784 KOps/s $\color{#35bf28}+1.90\%$
test_func_call_cm_runtime[False-eager] 0.9394ms 0.8439ms 1.1850 KOps/s 1.1187 KOps/s $\textbf{\color{#35bf28}+5.93\%}$
test_func_call_cm_runtime[False-compile] 1.4149ms 0.9030ms 1.1074 KOps/s 1.0931 KOps/s $\color{#35bf28}+1.31\%$
test_func_call_cm_runtime[False-compile-overhead] 0.5869ms 0.4612ms 2.1684 KOps/s 2.1380 KOps/s $\color{#35bf28}+1.42\%$
test_func_call_cm_runtime[True-eager] 1.6677ms 1.2268ms 815.1092 Ops/s 810.0017 Ops/s $\color{#35bf28}+0.63\%$
test_func_call_cm_runtime[True-compile] 1.4609ms 0.9508ms 1.0517 KOps/s 1.0432 KOps/s $\color{#35bf28}+0.82\%$
test_func_call_cm_runtime[True-compile-overhead] 0.9306ms 0.5056ms 1.9778 KOps/s 1.9406 KOps/s $\color{#35bf28}+1.92\%$
test_vmap_func_call_cm_runtime[eager] 2.8534ms 2.3665ms 422.5672 Ops/s 420.2384 Ops/s $\color{#35bf28}+0.55\%$
test_vmap_func_call_cm_runtime[compile] 1.5018ms 0.9642ms 1.0371 KOps/s 1.0217 KOps/s $\color{#35bf28}+1.51\%$
test_vmap_func_call_cm_runtime[compile-overhead] 0.9525ms 0.5123ms 1.9519 KOps/s 1.9160 KOps/s $\color{#35bf28}+1.87\%$
test_distributed 0.8077ms 0.1516ms 6.5947 KOps/s 6.3145 KOps/s $\color{#35bf28}+4.44\%$
test_tdmodule 58.8330μs 27.0116μs 37.0211 KOps/s 36.6070 KOps/s $\color{#35bf28}+1.13\%$
test_tdmodule_dispatch 74.7740μs 45.3273μs 22.0618 KOps/s 22.0490 KOps/s $\color{#35bf28}+0.06\%$
test_tdseq 47.9920μs 26.8664μs 37.2212 KOps/s 37.1741 KOps/s $\color{#35bf28}+0.13\%$
test_tdseq_dispatch 71.1130μs 47.3115μs 21.1365 KOps/s 20.8547 KOps/s $\color{#35bf28}+1.35\%$
test_instantiation_functorch 2.1824ms 2.0577ms 485.9694 Ops/s 477.4987 Ops/s $\color{#35bf28}+1.77\%$
test_exec_functorch 0.2343ms 0.1785ms 5.6033 KOps/s 5.5358 KOps/s $\color{#35bf28}+1.22\%$
test_exec_functional_call 0.2009ms 0.1608ms 6.2175 KOps/s 6.1548 KOps/s $\color{#35bf28}+1.02\%$
test_exec_td_decorator 0.4360ms 0.2327ms 4.2979 KOps/s 4.1967 KOps/s $\color{#35bf28}+2.41\%$
test_vmap_mlp_speed_decorator[True-True] 1.0156ms 0.8207ms 1.2185 KOps/s 1.2026 KOps/s $\color{#35bf28}+1.32\%$
test_vmap_mlp_speed_decorator[True-False] 0.9939ms 0.8192ms 1.2207 KOps/s 1.2063 KOps/s $\color{#35bf28}+1.19\%$
test_vmap_mlp_speed_decorator[False-True] 0.8585ms 0.7083ms 1.4117 KOps/s 1.3939 KOps/s $\color{#35bf28}+1.28\%$
test_vmap_mlp_speed_decorator[False-False] 0.8758ms 0.7079ms 1.4126 KOps/s 1.3967 KOps/s $\color{#35bf28}+1.14\%$
test_vmap_transformer_speed_decorator[True-True] 21.2746ms 20.4996ms 48.7815 Ops/s 48.6033 Ops/s $\color{#35bf28}+0.37\%$
test_vmap_transformer_speed_decorator[True-False] 21.1279ms 20.5152ms 48.7442 Ops/s 48.5652 Ops/s $\color{#35bf28}+0.37\%$
test_vmap_transformer_speed_decorator[False-True] 20.4066ms 20.2571ms 49.3654 Ops/s 49.0412 Ops/s $\color{#35bf28}+0.66\%$
test_vmap_transformer_speed_decorator[False-False] 20.3880ms 20.2743ms 49.3235 Ops/s 48.9395 Ops/s $\color{#35bf28}+0.78\%$
test_to_module_speed[True] 1.5710ms 1.4546ms 687.4934 Ops/s 677.9400 Ops/s $\color{#35bf28}+1.41\%$
test_to_module_speed[False] 1.5508ms 1.4438ms 692.5950 Ops/s 684.0056 Ops/s $\color{#35bf28}+1.26\%$
test_tc_init 76.9140μs 44.9789μs 22.2326 KOps/s 22.3555 KOps/s $\color{#d91a1a}-0.55\%$
test_tc_init_tensor_only 40.2320μs 9.7799μs 102.2503 KOps/s 102.4906 KOps/s $\color{#d91a1a}-0.23\%$
test_tc_init_nested 0.1254ms 88.8046μs 11.2607 KOps/s 11.2150 KOps/s $\color{#35bf28}+0.41\%$
test_tc_init_many_fields 44.6930μs 16.4716μs 60.7104 KOps/s 60.9511 KOps/s $\color{#d91a1a}-0.39\%$
test_tc_first_layer_tensor 32.9110μs 1.8123μs 551.7965 KOps/s 544.8713 KOps/s $\color{#35bf28}+1.27\%$
test_tc_first_layer_tensor_only 3.9502μs 0.4006μs 2.4965 MOps/s 2.4908 MOps/s $\color{#35bf28}+0.23\%$
test_tc_first_layer_tensor_set 30.7410μs 3.8996μs 256.4334 KOps/s 253.4918 KOps/s $\color{#35bf28}+1.16\%$
test_tc_first_layer_tensor_only_set 14.2700μs 3.4264μs 291.8533 KOps/s 304.8692 KOps/s $\color{#d91a1a}-4.27\%$
test_tc_first_layer_nontensor 26.6020μs 6.1508μs 162.5808 KOps/s 163.1925 KOps/s $\color{#d91a1a}-0.37\%$
test_tc_second_layer_tensor 35.6610μs 4.4213μs 226.1793 KOps/s 226.2436 KOps/s $\color{#d91a1a}-0.03\%$
test_tc_second_layer_nontensor 35.7420μs 8.7241μs 114.6254 KOps/s 115.5256 KOps/s $\color{#d91a1a}-0.78\%$
test_unbind 0.2525s 14.2693ms 70.0808 Ops/s 64.2115 Ops/s $\textbf{\color{#35bf28}+9.14\%}$
test_full_like 4.9620ms 4.3845ms 228.0778 Ops/s 59.6402 Ops/s $\textbf{\color{#35bf28}+282.42\%}$
test_zeros_like 10.4894ms 4.4497ms 224.7320 Ops/s 60.1633 Ops/s $\textbf{\color{#35bf28}+273.54\%}$
test_ones_like 4.8434ms 4.3837ms 228.1202 Ops/s 73.6575 Ops/s $\textbf{\color{#35bf28}+209.70\%}$
test_clone 6.8877ms 6.5724ms 152.1520 Ops/s 56.4040 Ops/s $\textbf{\color{#35bf28}+169.75\%}$
test_squeeze 0.1820ms 14.6876μs 68.0845 KOps/s 70.6794 KOps/s $\color{#d91a1a}-3.67\%$
test_unsqueeze 0.1643ms 0.1156ms 8.6525 KOps/s 8.8750 KOps/s $\color{#d91a1a}-2.51\%$
test_split 0.2400ms 0.1830ms 5.4655 KOps/s 5.3406 KOps/s $\color{#35bf28}+2.34\%$
test_permute 0.2747ms 0.2101ms 4.7595 KOps/s 4.8860 KOps/s $\color{#d91a1a}-2.59\%$
test_stack 51.7563ms 51.2776ms 19.5017 Ops/s 19.3409 Ops/s $\color{#35bf28}+0.83\%$
test_cat 51.4945ms 51.1515ms 19.5498 Ops/s 19.4404 Ops/s $\color{#35bf28}+0.56\%$
test_sequential_tensordict 0.2563ms 0.2128ms 4.6993 KOps/s 4.3667 KOps/s $\textbf{\color{#35bf28}+7.62\%}$
test_sequential_graph_module 0.5320ms 0.1191ms 8.3931 KOps/s 8.0173 KOps/s $\color{#35bf28}+4.69\%$
test_nested_tensordict 0.3579ms 0.2777ms 3.6004 KOps/s 3.3779 KOps/s $\textbf{\color{#35bf28}+6.59\%}$
test_nested_graph_module 0.5417ms 0.1291ms 7.7450 KOps/s 7.3726 KOps/s $\textbf{\color{#35bf28}+5.05\%}$

@github-actions

github-actions Bot commented Mar 10, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

$\color{#D29922}\textsf{\Large⚠\kern{0.2cm}\normalsize Warning}$ Result of GPU Benchmark Tests

Total Benchmarks: 261. Improved: $\large\color{#35bf28}12$. Worsened: $\large\color{#d91a1a}15$.

Expand to view detailed results
Name Max Mean Ops Ops on Repo HEAD Change
test_plain_set_nested 40.2610μs 14.3151μs 69.8562 KOps/s 69.8823 KOps/s $\color{#d91a1a}-0.04\%$
test_plain_set_stack_nested 42.1910μs 14.5686μs 68.6408 KOps/s 68.1611 KOps/s $\color{#35bf28}+0.70\%$
test_plain_set_nested_inplace 40.2800μs 16.2801μs 61.4248 KOps/s 62.0531 KOps/s $\color{#d91a1a}-1.01\%$
test_plain_set_stack_nested_inplace 40.2410μs 16.0672μs 62.2384 KOps/s 63.7998 KOps/s $\color{#d91a1a}-2.45\%$
test_items 31.1900μs 5.4206μs 184.4828 KOps/s 181.5541 KOps/s $\color{#35bf28}+1.61\%$
test_items_nested 0.4803ms 0.4513ms 2.2160 KOps/s 2.2518 KOps/s $\color{#d91a1a}-1.59\%$
test_items_nested_locked 0.5222ms 0.4528ms 2.2083 KOps/s 2.2425 KOps/s $\color{#d91a1a}-1.53\%$
test_items_nested_leaf 0.1532ms 93.0038μs 10.7523 KOps/s 10.7900 KOps/s $\color{#d91a1a}-0.35\%$
test_items_stack_nested 0.4802ms 0.4505ms 2.2198 KOps/s 2.2367 KOps/s $\color{#d91a1a}-0.76\%$
test_items_stack_nested_leaf 0.1224ms 92.9423μs 10.7594 KOps/s 10.7612 KOps/s $\color{#d91a1a}-0.02\%$
test_items_stack_nested_locked 0.4825ms 0.4521ms 2.2117 KOps/s 2.2345 KOps/s $\color{#d91a1a}-1.02\%$
test_keys 26.0100μs 4.1225μs 242.5694 KOps/s 241.6150 KOps/s $\color{#35bf28}+0.39\%$
test_keys_nested 0.1549ms 0.1285ms 7.7821 KOps/s 7.8446 KOps/s $\color{#d91a1a}-0.80\%$
test_keys_nested_locked 0.7876ms 0.1355ms 7.3777 KOps/s 7.3866 KOps/s $\color{#d91a1a}-0.12\%$
test_keys_nested_leaf 0.1478ms 0.1185ms 8.4418 KOps/s 8.4847 KOps/s $\color{#d91a1a}-0.51\%$
test_keys_stack_nested 0.1589ms 0.1275ms 7.8415 KOps/s 7.8388 KOps/s $\color{#35bf28}+0.03\%$
test_keys_stack_nested_leaf 0.1665ms 0.1182ms 8.4588 KOps/s 8.4949 KOps/s $\color{#d91a1a}-0.42\%$
test_keys_stack_nested_locked 0.1920ms 0.1356ms 7.3766 KOps/s 7.4643 KOps/s $\color{#d91a1a}-1.17\%$
test_values 5.2360μs 1.0103μs 989.7932 KOps/s 1.0007 MOps/s $\color{#d91a1a}-1.09\%$
test_values_nested 76.1210μs 51.7192μs 19.3352 KOps/s 19.5554 KOps/s $\color{#d91a1a}-1.13\%$
test_values_nested_locked 88.3010μs 54.7508μs 18.2646 KOps/s 18.4323 KOps/s $\color{#d91a1a}-0.91\%$
test_values_nested_leaf 0.1048ms 59.2672μs 16.8728 KOps/s 17.0631 KOps/s $\color{#d91a1a}-1.12\%$
test_values_stack_nested 84.6610μs 51.7659μs 19.3177 KOps/s 19.3997 KOps/s $\color{#d91a1a}-0.42\%$
test_values_stack_nested_leaf 99.0910μs 59.5036μs 16.8057 KOps/s 17.0760 KOps/s $\color{#d91a1a}-1.58\%$
test_values_stack_nested_locked 84.1410μs 55.0235μs 18.1741 KOps/s 18.2332 KOps/s $\color{#d91a1a}-0.32\%$
test_membership 5.0152μs 0.8162μs 1.2252 MOps/s 1.2114 MOps/s $\color{#35bf28}+1.14\%$
test_membership_nested 30.1700μs 2.7426μs 364.6238 KOps/s 362.8564 KOps/s $\color{#35bf28}+0.49\%$
test_membership_nested_leaf 16.8305μs 2.6699μs 374.5455 KOps/s 365.1363 KOps/s $\color{#35bf28}+2.58\%$
test_membership_stacked_nested 22.5700μs 2.7651μs 361.6552 KOps/s 364.1107 KOps/s $\color{#d91a1a}-0.67\%$
test_membership_stacked_nested_leaf 36.4410μs 2.7343μs 365.7238 KOps/s 363.7383 KOps/s $\color{#35bf28}+0.55\%$
test_membership_nested_last 31.9700μs 4.1545μs 240.7012 KOps/s 242.1245 KOps/s $\color{#d91a1a}-0.59\%$
test_membership_nested_leaf_last 33.5110μs 4.1231μs 242.5372 KOps/s 242.2901 KOps/s $\color{#35bf28}+0.10\%$
test_membership_stacked_nested_last 37.5810μs 4.1144μs 243.0504 KOps/s 242.9470 KOps/s $\color{#35bf28}+0.04\%$
test_membership_stacked_nested_leaf_last 28.0800μs 4.1349μs 241.8451 KOps/s 244.5513 KOps/s $\color{#d91a1a}-1.11\%$
test_nested_getleaf 54.7600μs 20.6663μs 48.3880 KOps/s 49.8407 KOps/s $\color{#d91a1a}-2.91\%$
test_nested_get 52.3010μs 19.6577μs 50.8707 KOps/s 51.9714 KOps/s $\color{#d91a1a}-2.12\%$
test_stacked_getleaf 45.7410μs 20.6325μs 48.4672 KOps/s 49.4609 KOps/s $\color{#d91a1a}-2.01\%$
test_stacked_get 0.1065ms 19.6146μs 50.9824 KOps/s 52.8046 KOps/s $\color{#d91a1a}-3.45\%$
test_nested_getitemleaf 51.8810μs 21.2260μs 47.1120 KOps/s 47.9492 KOps/s $\color{#d91a1a}-1.75\%$
test_nested_getitem 46.6400μs 20.1541μs 49.6177 KOps/s 50.9639 KOps/s $\color{#d91a1a}-2.64\%$
test_stacked_getitemleaf 46.3900μs 21.0057μs 47.6062 KOps/s 48.1376 KOps/s $\color{#d91a1a}-1.10\%$
test_stacked_getitem 49.8210μs 20.1126μs 49.7200 KOps/s 50.8546 KOps/s $\color{#d91a1a}-2.23\%$
test_lock_nested 8.2552ms 0.4634ms 2.1579 KOps/s 2.2058 KOps/s $\color{#d91a1a}-2.17\%$
test_lock_stack_nested 0.4919ms 0.4552ms 2.1968 KOps/s 2.1698 KOps/s $\color{#35bf28}+1.25\%$
test_unlock_nested 0.4849ms 0.3719ms 2.6890 KOps/s 2.7033 KOps/s $\color{#d91a1a}-0.53\%$
test_unlock_stack_nested 0.4753ms 0.3675ms 2.7211 KOps/s 2.6764 KOps/s $\color{#35bf28}+1.67\%$
test_flatten_speed 0.1727ms 0.1173ms 8.5239 KOps/s 8.4965 KOps/s $\color{#35bf28}+0.32\%$
test_unflatten_speed 0.6597ms 0.5506ms 1.8163 KOps/s 1.8420 KOps/s $\color{#d91a1a}-1.39\%$
test_common_ops 0.8175ms 0.6726ms 1.4868 KOps/s 1.4721 KOps/s $\color{#35bf28}+1.00\%$
test_creation 71.4810μs 2.8400μs 352.1135 KOps/s 338.9679 KOps/s $\color{#35bf28}+3.88\%$
test_creation_empty 34.5810μs 6.5098μs 153.6157 KOps/s 151.8745 KOps/s $\color{#35bf28}+1.15\%$
test_creation_nested_1 31.9310μs 10.7925μs 92.6569 KOps/s 90.2441 KOps/s $\color{#35bf28}+2.67\%$
test_creation_nested_2 43.2200μs 12.4650μs 80.2247 KOps/s 79.4470 KOps/s $\color{#35bf28}+0.98\%$
test_creation_many_keys[10] 46.4800μs 19.5543μs 51.1397 KOps/s 51.1186 KOps/s $\color{#35bf28}+0.04\%$
test_creation_many_keys[50] 0.1108ms 84.2405μs 11.8708 KOps/s 11.7211 KOps/s $\color{#35bf28}+1.28\%$
test_creation_many_keys[100] 0.1890ms 0.1650ms 6.0591 KOps/s 5.9629 KOps/s $\color{#35bf28}+1.61\%$
test_creation_nested_many_keys[10] 75.6810μs 42.2519μs 23.6675 KOps/s 23.6042 KOps/s $\color{#35bf28}+0.27\%$
test_creation_nested_many_keys[50] 0.2237ms 0.1712ms 5.8395 KOps/s 5.7598 KOps/s $\color{#35bf28}+1.38\%$
test_clone 60.5200μs 12.6589μs 78.9955 KOps/s 78.1500 KOps/s $\color{#35bf28}+1.08\%$
test_getitem[int] 1.6471ms 14.5161μs 68.8892 KOps/s 62.4082 KOps/s $\textbf{\color{#35bf28}+10.38\%}$
test_getitem[slice_int] 0.1370ms 23.1284μs 43.2369 KOps/s 42.8124 KOps/s $\color{#35bf28}+0.99\%$
test_getitem[range] 0.1767ms 60.5721μs 16.5092 KOps/s 16.4154 KOps/s $\color{#35bf28}+0.57\%$
test_getitem[tuple] 0.1330ms 22.9671μs 43.5406 KOps/s 43.6205 KOps/s $\color{#d91a1a}-0.18\%$
test_getitem[list] 0.1797ms 56.2842μs 17.7670 KOps/s 17.0817 KOps/s $\color{#35bf28}+4.01\%$
test_setitem_dim[int] 51.4410μs 25.1113μs 39.8228 KOps/s 38.5231 KOps/s $\color{#35bf28}+3.37\%$
test_setitem_dim[slice_int] 88.5810μs 42.0029μs 23.8079 KOps/s 23.0603 KOps/s $\color{#35bf28}+3.24\%$
test_setitem_dim[range] 0.1165ms 90.9266μs 10.9979 KOps/s 10.4635 KOps/s $\textbf{\color{#35bf28}+5.11\%}$
test_setitem_dim[tuple] 58.3910μs 37.4649μs 26.6916 KOps/s 25.7565 KOps/s $\color{#35bf28}+3.63\%$
test_setitem 49.5510μs 16.8500μs 59.3472 KOps/s 58.4898 KOps/s $\color{#35bf28}+1.47\%$
test_set 53.0900μs 16.0960μs 62.1272 KOps/s 61.4713 KOps/s $\color{#35bf28}+1.07\%$
test_set_shared 0.5478ms 0.2008ms 4.9799 KOps/s 4.9432 KOps/s $\color{#35bf28}+0.74\%$
test_update 92.8110μs 20.4246μs 48.9606 KOps/s 47.6770 KOps/s $\color{#35bf28}+2.69\%$
test_update_nested 89.1210μs 31.1256μs 32.1279 KOps/s 31.0704 KOps/s $\color{#35bf28}+3.40\%$
test_update__nested 0.5497ms 32.9427μs 30.3557 KOps/s 30.1275 KOps/s $\color{#35bf28}+0.76\%$
test_set_nested 47.4000μs 17.7356μs 56.3838 KOps/s 54.8675 KOps/s $\color{#35bf28}+2.76\%$
test_set_nested_new 95.0510μs 22.4287μs 44.5856 KOps/s 43.1377 KOps/s $\color{#35bf28}+3.36\%$
test_select 79.3010μs 39.6765μs 25.2039 KOps/s 25.4601 KOps/s $\color{#d91a1a}-1.01\%$
test_select_nested 0.1051ms 70.6319μs 14.1579 KOps/s 14.1487 KOps/s $\color{#35bf28}+0.07\%$
test_exclude_nested 0.1203ms 86.5075μs 11.5597 KOps/s 11.4775 KOps/s $\color{#35bf28}+0.72\%$
test_empty[True] 0.7628ms 0.3863ms 2.5883 KOps/s 2.6056 KOps/s $\color{#d91a1a}-0.66\%$
test_empty[False] 7.6400μs 1.2579μs 794.9857 KOps/s 806.6480 KOps/s $\color{#d91a1a}-1.45\%$
test_to 0.1080ms 71.9359μs 13.9013 KOps/s 13.9504 KOps/s $\color{#d91a1a}-0.35\%$
test_to_nonblocking 0.1032ms 64.7980μs 15.4326 KOps/s 15.9820 KOps/s $\color{#d91a1a}-3.44\%$
test_unbind_speed 0.3782ms 0.3117ms 3.2083 KOps/s 3.1595 KOps/s $\color{#35bf28}+1.55\%$
test_unbind_speed_stack0 0.4164ms 0.3190ms 3.1348 KOps/s 3.1492 KOps/s $\color{#d91a1a}-0.46\%$
test_unbind_speed_stack1 0.1048s 0.8059ms 1.2409 KOps/s 1.2300 KOps/s $\color{#35bf28}+0.88\%$
test_split 0.1044s 1.2042ms 830.4085 Ops/s 818.4326 Ops/s $\color{#35bf28}+1.46\%$
test_chunk 0.1050s 1.1613ms 861.1141 Ops/s 956.6955 Ops/s $\textbf{\color{#d91a1a}-9.99\%}$
test_to_cpu_blocking 19.7017ms 19.3320ms 51.7277 Ops/s 46.2910 Ops/s $\textbf{\color{#35bf28}+11.74\%}$
test_to_cpu_global_sync 11.3304ms 11.2488ms 88.8984 Ops/s 88.8262 Ops/s $\color{#35bf28}+0.08\%$
test_to_cpu_event_sync 12.5581ms 12.2306ms 81.7623 Ops/s 81.4836 Ops/s $\color{#35bf28}+0.34\%$
test_to_cpu_default 0.1172s 13.5391ms 73.8601 Ops/s 81.2872 Ops/s $\textbf{\color{#d91a1a}-9.14\%}$
test_consolidate[False-None] 4.0741ms 3.9709ms 251.8306 Ops/s 224.7073 Ops/s $\textbf{\color{#35bf28}+12.07\%}$
test_consolidate[default-None] 2.2885ms 1.9677ms 508.2119 Ops/s 495.2137 Ops/s $\color{#35bf28}+2.62\%$
test_consolidate[reduce-overhead-None] 2.0623ms 1.8941ms 527.9627 Ops/s 516.9668 Ops/s $\color{#35bf28}+2.13\%$
test_consolidate_njt[False-None] 8.4526ms 8.1758ms 122.3123 Ops/s 123.9258 Ops/s $\color{#d91a1a}-1.30\%$
test_to[False-False-None] 2.2620ms 2.0730ms 482.3862 Ops/s 491.3994 Ops/s $\color{#d91a1a}-1.83\%$
test_to[True-False-None] 2.0684ms 1.8759ms 533.0830 Ops/s 546.8768 Ops/s $\color{#d91a1a}-2.52\%$
test_to[within-False-None] 6.3075ms 6.0058ms 166.5066 Ops/s 170.1851 Ops/s $\color{#d91a1a}-2.16\%$
test_to[True-default-None] 8.6736ms 8.5086ms 117.5276 Ops/s 118.5202 Ops/s $\color{#d91a1a}-0.84\%$
test_to_njt[False-False-None] 8.4703ms 8.2378ms 121.3917 Ops/s 121.5975 Ops/s $\color{#d91a1a}-0.17\%$
test_to_njt[True-False-None] 7.1008ms 6.6839ms 149.6127 Ops/s 150.4387 Ops/s $\color{#d91a1a}-0.55\%$
test_to_njt[within-False-None] 15.1231ms 14.9081ms 67.0776 Ops/s 67.4336 Ops/s $\color{#d91a1a}-0.53\%$
test_creation[device0] 0.3894ms 0.1127ms 8.8718 KOps/s 8.7632 KOps/s $\color{#35bf28}+1.24\%$
test_creation_from_tensor 0.4469ms 0.1101ms 9.0851 KOps/s 9.0962 KOps/s $\color{#d91a1a}-0.12\%$
test_add_one[memmap_tensor0] 0.2594ms 6.3878μs 156.5479 KOps/s 157.3759 KOps/s $\color{#d91a1a}-0.53\%$
test_contiguous[memmap_tensor0] 26.2710μs 0.6132μs 1.6307 MOps/s 2.2527 MOps/s $\textbf{\color{#d91a1a}-27.61\%}$
test_stack[memmap_tensor0] 41.8800μs 4.6624μs 214.4809 KOps/s 216.4700 KOps/s $\color{#d91a1a}-0.92\%$
test_memmaptd_index 1.0818ms 0.2589ms 3.8624 KOps/s 3.9059 KOps/s $\color{#d91a1a}-1.11\%$
test_memmaptd_index_astensor 0.5058ms 0.3567ms 2.8033 KOps/s 2.8344 KOps/s $\color{#d91a1a}-1.10\%$
test_memmaptd_index_op 0.9606ms 0.5978ms 1.6729 KOps/s 1.6857 KOps/s $\color{#d91a1a}-0.76\%$
test_serialize_model 0.1360s 0.1343s 7.4470 Ops/s 7.4349 Ops/s $\color{#35bf28}+0.16\%$
test_serialize_model_pickle 1.3785s 1.2004s 0.8330 Ops/s 0.8385 Ops/s $\color{#d91a1a}-0.65\%$
test_serialize_weights 0.1382s 0.1341s 7.4591 Ops/s 7.4479 Ops/s $\color{#35bf28}+0.15\%$
test_serialize_weights_returnearly 0.3489s 84.6966ms 11.8068 Ops/s 15.5980 Ops/s $\textbf{\color{#d91a1a}-24.31\%}$
test_serialize_weights_pickle 1.3586s 1.2124s 0.8248 Ops/s 0.8215 Ops/s $\color{#35bf28}+0.41\%$
test_reshape_pytree 0.2238ms 31.4475μs 31.7990 KOps/s 32.0803 KOps/s $\color{#d91a1a}-0.88\%$
test_reshape_td 80.1710μs 42.5626μs 23.4948 KOps/s 23.2810 KOps/s $\color{#35bf28}+0.92\%$
test_view_pytree 0.2150ms 31.3446μs 31.9034 KOps/s 31.9393 KOps/s $\color{#d91a1a}-0.11\%$
test_view_td 83.0010μs 51.4763μs 19.4264 KOps/s 19.9784 KOps/s $\color{#d91a1a}-2.76\%$
test_unbind_pytree 0.2399ms 35.1467μs 28.4521 KOps/s 28.4388 KOps/s $\color{#35bf28}+0.05\%$
test_unbind_td 0.1383ms 46.5950μs 21.4615 KOps/s 21.2808 KOps/s $\color{#35bf28}+0.85\%$
test_split_pytree 0.2420ms 40.3057μs 24.8104 KOps/s 25.0341 KOps/s $\color{#d91a1a}-0.89\%$
test_split_td 0.1462ms 61.3156μs 16.3091 KOps/s 16.5131 KOps/s $\color{#d91a1a}-1.24\%$
test_add_pytree 0.2263ms 40.6931μs 24.5742 KOps/s 25.3462 KOps/s $\color{#d91a1a}-3.05\%$
test_add_td 92.4820μs 52.2201μs 19.1497 KOps/s 19.4036 KOps/s $\color{#d91a1a}-1.31\%$
test_compile_add_one_nested[tensordict-compile] 0.3071ms 0.1408ms 7.1043 KOps/s 6.8603 KOps/s $\color{#35bf28}+3.56\%$
test_compile_add_one_nested[tensordict-eager] 0.5505ms 0.1943ms 5.1480 KOps/s 5.3093 KOps/s $\color{#d91a1a}-3.04\%$
test_compile_add_one_nested[pytree-compile] 0.2624ms 0.1039ms 9.6243 KOps/s 9.3122 KOps/s $\color{#35bf28}+3.35\%$
test_compile_add_one_nested[pytree-eager] 0.6039ms 0.1718ms 5.8221 KOps/s 5.5751 KOps/s $\color{#35bf28}+4.43\%$
test_compile_copy_nested[tensordict-compile] 0.3668ms 9.9929μs 100.0707 KOps/s 99.8980 KOps/s $\color{#35bf28}+0.17\%$
test_compile_copy_nested[tensordict-eager] 0.1008ms 51.2450μs 19.5141 KOps/s 19.5559 KOps/s $\color{#d91a1a}-0.21\%$
test_compile_copy_nested[pytree-compile] 41.7600μs 9.6176μs 103.9756 KOps/s 106.8039 KOps/s $\color{#d91a1a}-2.65\%$
test_compile_copy_nested[pytree-eager] 0.4520ms 65.3464μs 15.3031 KOps/s 15.1308 KOps/s $\color{#35bf28}+1.14\%$
test_compile_add_one_flat[tensordict-compile] 0.2274ms 0.1727ms 5.7917 KOps/s 5.4981 KOps/s $\textbf{\color{#35bf28}+5.34\%}$
test_compile_add_one_flat[tensordict-eager] 0.3812ms 0.2739ms 3.6514 KOps/s 3.6524 KOps/s $\color{#d91a1a}-0.03\%$
test_compile_add_one_flat[tensorclass-compile] 0.1910ms 0.1140ms 8.7757 KOps/s 8.4859 KOps/s $\color{#35bf28}+3.41\%$
test_compile_add_one_flat[tensorclass-eager] 0.1187ms 71.3358μs 14.0182 KOps/s 13.8009 KOps/s $\color{#35bf28}+1.57\%$
test_compile_add_one_flat[pytree-compile] 0.2240ms 0.1548ms 6.4598 KOps/s 6.3011 KOps/s $\color{#35bf28}+2.52\%$
test_compile_add_one_flat[pytree-eager] 0.8106ms 0.5080ms 1.9684 KOps/s 1.9348 KOps/s $\color{#35bf28}+1.74\%$
test_compile_add_self_flat[tensordict-eager] 0.3635ms 0.3258ms 3.0694 KOps/s 3.0584 KOps/s $\color{#35bf28}+0.36\%$
test_compile_add_self_flat[tensordict-compile] 0.2156ms 0.1765ms 5.6660 KOps/s 5.3984 KOps/s $\color{#35bf28}+4.96\%$
test_compile_add_self_flat[tensorclass-eager] 0.1305ms 87.0782μs 11.4839 KOps/s 11.2209 KOps/s $\color{#35bf28}+2.34\%$
test_compile_add_self_flat[tensorclass-compile] 0.1691ms 0.1185ms 8.4366 KOps/s 7.7545 KOps/s $\textbf{\color{#35bf28}+8.80\%}$
test_compile_add_self_flat[pytree-eager] 0.6973ms 0.4686ms 2.1340 KOps/s 2.2861 KOps/s $\textbf{\color{#d91a1a}-6.65\%}$
test_compile_add_self_flat[pytree-compile] 0.2204ms 0.1554ms 6.4360 KOps/s 6.2120 KOps/s $\color{#35bf28}+3.61\%$
test_compile_copy_flat[tensordict-compile] 50.7010μs 12.9769μs 77.0602 KOps/s 75.2555 KOps/s $\color{#35bf28}+2.40\%$
test_compile_copy_flat[tensordict-eager] 75.4910μs 40.2418μs 24.8498 KOps/s 25.1554 KOps/s $\color{#d91a1a}-1.22\%$
test_compile_copy_flat[pytree-compile] 39.2700μs 10.2312μs 97.7399 KOps/s 95.2660 KOps/s $\color{#35bf28}+2.60\%$
test_compile_copy_flat[pytree-eager] 0.4051ms 51.0119μs 19.6033 KOps/s 19.6497 KOps/s $\color{#d91a1a}-0.24\%$
test_compile_assign_and_add[tensordict-compile] 1.9438ms 0.1695ms 5.8997 KOps/s 5.5632 KOps/s $\textbf{\color{#35bf28}+6.05\%}$
test_compile_assign_and_add[tensordict-eager] 3.3427ms 3.2556ms 307.1593 Ops/s 317.0908 Ops/s $\color{#d91a1a}-3.13\%$
test_compile_assign_and_add[pytree-compile] 1.9371ms 0.1589ms 6.2927 KOps/s 6.2545 KOps/s $\color{#35bf28}+0.61\%$
test_compile_assign_and_add[pytree-eager] 3.0140ms 2.7445ms 364.3599 Ops/s 377.6607 Ops/s $\color{#d91a1a}-3.52\%$
test_compile_indexing[tensor-tensordict-compile] 0.1641ms 0.1102ms 9.0781 KOps/s 8.5067 KOps/s $\textbf{\color{#35bf28}+6.72\%}$
test_compile_indexing[tensor-tensordict-eager] 0.3124ms 75.0442μs 13.3255 KOps/s 13.4515 KOps/s $\color{#d91a1a}-0.94\%$
test_compile_indexing[tensor-tensorclass-compile] 0.1410ms 94.6766μs 10.5623 KOps/s 10.5973 KOps/s $\color{#d91a1a}-0.33\%$
test_compile_indexing[tensor-tensorclass-eager] 0.2626ms 45.7180μs 21.8732 KOps/s 23.2721 KOps/s $\textbf{\color{#d91a1a}-6.01\%}$
test_compile_indexing[tensor-pytree-compile] 0.1408ms 96.1208μs 10.4036 KOps/s 10.4756 KOps/s $\color{#d91a1a}-0.69\%$
test_compile_indexing[tensor-pytree-eager] 0.2357ms 45.8901μs 21.7912 KOps/s 23.1333 KOps/s $\textbf{\color{#d91a1a}-5.80\%}$
test_compile_indexing[slice-tensordict-compile] 93.1110μs 55.0767μs 18.1565 KOps/s 17.2644 KOps/s $\textbf{\color{#35bf28}+5.17\%}$
test_compile_indexing[slice-tensordict-eager] 0.2138ms 25.7440μs 38.8440 KOps/s 37.5858 KOps/s $\color{#35bf28}+3.35\%$
test_compile_indexing[slice-tensorclass-compile] 93.8910μs 42.7167μs 23.4100 KOps/s 22.9586 KOps/s $\color{#35bf28}+1.97\%$
test_compile_indexing[slice-tensorclass-eager] 0.2469ms 21.2961μs 46.9570 KOps/s 45.8817 KOps/s $\color{#35bf28}+2.34\%$
test_compile_indexing[slice-pytree-compile] 79.3110μs 43.0347μs 23.2371 KOps/s 22.5258 KOps/s $\color{#35bf28}+3.16\%$
test_compile_indexing[slice-pytree-eager] 0.2586ms 21.0440μs 47.5194 KOps/s 45.9500 KOps/s $\color{#35bf28}+3.42\%$
test_compile_indexing[int-tensordict-compile] 94.3720μs 54.8285μs 18.2387 KOps/s 17.3942 KOps/s $\color{#35bf28}+4.86\%$
test_compile_indexing[int-tensordict-eager] 0.2215ms 25.4765μs 39.2519 KOps/s 37.9938 KOps/s $\color{#35bf28}+3.31\%$
test_compile_indexing[int-tensorclass-compile] 85.0510μs 42.8099μs 23.3591 KOps/s 23.1642 KOps/s $\color{#35bf28}+0.84\%$
test_compile_indexing[int-tensorclass-eager] 0.2448ms 21.2072μs 47.1538 KOps/s 46.0010 KOps/s $\color{#35bf28}+2.51\%$
test_compile_indexing[int-pytree-compile] 83.7410μs 43.2789μs 23.1059 KOps/s 22.6231 KOps/s $\color{#35bf28}+2.13\%$
test_compile_indexing[int-pytree-eager] 0.3049ms 21.3808μs 46.7709 KOps/s 46.0472 KOps/s $\color{#35bf28}+1.57\%$
test_compile_replace[single-eager] 0.1081ms 48.1698μs 20.7599 KOps/s 22.0912 KOps/s $\textbf{\color{#d91a1a}-6.03\%}$
test_compile_replace[single-compile] 0.1475ms 0.1021ms 9.7964 KOps/s 9.3424 KOps/s $\color{#35bf28}+4.86\%$
test_compile_replace[multi-eager] 0.6572ms 0.5770ms 1.7330 KOps/s 1.8275 KOps/s $\textbf{\color{#d91a1a}-5.17\%}$
test_compile_replace[multi-compile] 0.1580ms 0.1093ms 9.1461 KOps/s 8.9864 KOps/s $\color{#35bf28}+1.78\%$
test_compile_tc_getattr_20[eager] 0.2401ms 0.1840ms 5.4352 KOps/s 5.9697 KOps/s $\textbf{\color{#d91a1a}-8.95\%}$
test_compile_tc_getattr_20[compile] 0.1659ms 0.1168ms 8.5607 KOps/s 8.4187 KOps/s $\color{#35bf28}+1.69\%$
test_compile_clone_shallow[20-eager] 50.2610μs 18.4601μs 54.1708 KOps/s 53.1558 KOps/s $\color{#35bf28}+1.91\%$
test_compile_clone_shallow[20-compile] 45.2700μs 10.8732μs 91.9690 KOps/s 87.5862 KOps/s $\textbf{\color{#35bf28}+5.00\%}$
test_compile_clone_shallow[40-eager] 64.4010μs 32.3808μs 30.8825 KOps/s 30.3505 KOps/s $\color{#35bf28}+1.75\%$
test_compile_clone_shallow[40-compile] 50.2800μs 11.9448μs 83.7185 KOps/s 80.5668 KOps/s $\color{#35bf28}+3.91\%$
test_compile_clone_shallow[80-eager] 96.9410μs 60.7511μs 16.4606 KOps/s 16.3020 KOps/s $\color{#35bf28}+0.97\%$
test_compile_clone_shallow[80-compile] 48.3800μs 14.3666μs 69.6060 KOps/s 66.1862 KOps/s $\textbf{\color{#35bf28}+5.17\%}$
test_compile_update_inplace[eager] 90.0810μs 56.6285μs 17.6590 KOps/s 17.3332 KOps/s $\color{#35bf28}+1.88\%$
test_compile_update_inplace[compile] 0.1888ms 0.1345ms 7.4359 KOps/s 7.1376 KOps/s $\color{#35bf28}+4.18\%$
test_mod_add[eager] 0.1030ms 51.0790μs 19.5775 KOps/s 21.1419 KOps/s $\textbf{\color{#d91a1a}-7.40\%}$
test_mod_add[compile] 0.1865ms 0.1020ms 9.8010 KOps/s 9.5753 KOps/s $\color{#35bf28}+2.36\%$
test_mod_add[compile-overhead] 0.2492ms 0.1455ms 6.8723 KOps/s 6.6375 KOps/s $\color{#35bf28}+3.54\%$
test_mod_wrap[eager] 0.3689ms 0.2865ms 3.4906 KOps/s 3.5530 KOps/s $\color{#d91a1a}-1.76\%$
test_mod_wrap[compile] 0.4817ms 0.3352ms 2.9832 KOps/s 2.9053 KOps/s $\color{#35bf28}+2.68\%$
test_mod_wrap[compile-overhead] 7.4062ms 4.0418ms 247.4131 Ops/s 252.9096 Ops/s $\color{#d91a1a}-2.17\%$
test_mod_wrap_and_backward[eager] 1.6277ms 1.4677ms 681.3350 Ops/s 683.5161 Ops/s $\color{#d91a1a}-0.32\%$
test_mod_wrap_and_backward[compile] 1.4838ms 1.4004ms 714.1029 Ops/s 701.5052 Ops/s $\color{#35bf28}+1.80\%$
test_mod_wrap_and_backward[compile-overhead] 1.2172ms 0.8674ms 1.1528 KOps/s 1.1385 KOps/s $\color{#35bf28}+1.26\%$
test_seq_add[eager] 0.2795ms 0.1486ms 6.7315 KOps/s 6.7011 KOps/s $\color{#35bf28}+0.45\%$
test_seq_add[compile] 0.6374ms 0.1125ms 8.8891 KOps/s 8.7173 KOps/s $\color{#35bf28}+1.97\%$
test_seq_add[compile-overhead] 0.6111ms 0.1521ms 6.5747 KOps/s 6.4397 KOps/s $\color{#35bf28}+2.10\%$
test_seq_wrap[eager] 1.0016ms 0.5278ms 1.8945 KOps/s 1.9315 KOps/s $\color{#d91a1a}-1.92\%$
test_seq_wrap[compile] 0.8120ms 0.3675ms 2.7212 KOps/s 2.7874 KOps/s $\color{#d91a1a}-2.37\%$
test_seq_wrap[compile-overhead] 0.4034ms 0.2614ms 3.8256 KOps/s 3.8377 KOps/s $\color{#d91a1a}-0.31\%$
test_func_call_runtime[False-eager] 1.3103ms 0.8608ms 1.1617 KOps/s 1.2423 KOps/s $\textbf{\color{#d91a1a}-6.48\%}$
test_func_call_runtime[False-compile] 1.4070ms 0.9025ms 1.1081 KOps/s 1.1470 KOps/s $\color{#d91a1a}-3.39\%$
test_func_call_runtime[False-compile-overhead] 0.8837ms 0.4386ms 2.2798 KOps/s 2.2635 KOps/s $\color{#35bf28}+0.72\%$
test_func_call_runtime[True-eager] 1.5411ms 1.1010ms 908.2467 Ops/s 971.9069 Ops/s $\textbf{\color{#d91a1a}-6.55\%}$
test_func_call_runtime[True-compile] 1.4220ms 0.9129ms 1.0954 KOps/s 1.1323 KOps/s $\color{#d91a1a}-3.26\%$
test_func_call_runtime[True-compile-overhead] 0.9179ms 0.4600ms 2.1738 KOps/s 2.1826 KOps/s $\color{#d91a1a}-0.40\%$
test_func_call_cm_runtime[False-eager] 1.3164ms 0.8460ms 1.1820 KOps/s 1.2417 KOps/s $\color{#d91a1a}-4.81\%$
test_func_call_cm_runtime[False-compile] 1.0307ms 0.8855ms 1.1293 KOps/s 1.1490 KOps/s $\color{#d91a1a}-1.71\%$
test_func_call_cm_runtime[False-compile-overhead] 0.8694ms 0.4391ms 2.2775 KOps/s 2.2578 KOps/s $\color{#35bf28}+0.87\%$
test_func_call_cm_runtime[True-eager] 1.6365ms 1.2028ms 831.3773 Ops/s 851.4096 Ops/s $\color{#d91a1a}-2.35\%$
test_func_call_cm_runtime[True-compile] 1.1459ms 0.9134ms 1.0948 KOps/s 1.0895 KOps/s $\color{#35bf28}+0.48\%$
test_func_call_cm_runtime[True-compile-overhead] 0.9252ms 0.4854ms 2.0604 KOps/s 2.0356 KOps/s $\color{#35bf28}+1.22\%$
test_vmap_func_call_cm_runtime[eager] 2.8220ms 2.3216ms 430.7413 Ops/s 435.0886 Ops/s $\color{#d91a1a}-1.00\%$
test_vmap_func_call_cm_runtime[compile] 1.4112ms 0.9490ms 1.0537 KOps/s 1.0646 KOps/s $\color{#d91a1a}-1.02\%$
test_vmap_func_call_cm_runtime[compile-overhead] 0.9345ms 0.4919ms 2.0328 KOps/s 1.9957 KOps/s $\color{#35bf28}+1.86\%$
test_distributed 0.6559ms 0.1508ms 6.6334 KOps/s 6.5361 KOps/s $\color{#35bf28}+1.49\%$
test_tdmodule 58.4000μs 25.5755μs 39.1000 KOps/s 37.1809 KOps/s $\textbf{\color{#35bf28}+5.16\%}$
test_tdmodule_dispatch 71.6800μs 42.6393μs 23.4526 KOps/s 23.1429 KOps/s $\color{#35bf28}+1.34\%$
test_tdseq 55.8810μs 25.5516μs 39.1365 KOps/s 39.3445 KOps/s $\color{#d91a1a}-0.53\%$
test_tdseq_dispatch 65.2110μs 44.9734μs 22.2354 KOps/s 22.0600 KOps/s $\color{#35bf28}+0.79\%$
test_instantiation_functorch 2.3646ms 1.9435ms 514.5421 Ops/s 510.8197 Ops/s $\color{#35bf28}+0.73\%$
test_exec_functorch 0.2974ms 0.1727ms 5.7897 KOps/s 5.9506 KOps/s $\color{#d91a1a}-2.70\%$
test_exec_functional_call 0.5754ms 0.1527ms 6.5489 KOps/s 6.7104 KOps/s $\color{#d91a1a}-2.41\%$
test_exec_td_decorator 0.6456ms 0.2270ms 4.4062 KOps/s 4.4845 KOps/s $\color{#d91a1a}-1.75\%$
test_vmap_mlp_speed_decorator[True-True] 1.2855ms 0.8256ms 1.2113 KOps/s 1.2490 KOps/s $\color{#d91a1a}-3.02\%$
test_vmap_mlp_speed_decorator[True-False] 1.2538ms 0.8279ms 1.2079 KOps/s 1.2498 KOps/s $\color{#d91a1a}-3.35\%$
test_vmap_mlp_speed_decorator[False-True] 1.1649ms 0.7170ms 1.3946 KOps/s 1.4451 KOps/s $\color{#d91a1a}-3.49\%$
test_vmap_mlp_speed_decorator[False-False] 1.1609ms 0.7160ms 1.3967 KOps/s 1.4474 KOps/s $\color{#d91a1a}-3.50\%$
test_vmap_transformer_speed_decorator[True-True] 21.5062ms 20.6242ms 48.4867 Ops/s 49.8087 Ops/s $\color{#d91a1a}-2.65\%$
test_vmap_transformer_speed_decorator[True-False] 21.3734ms 20.6437ms 48.4410 Ops/s 49.8906 Ops/s $\color{#d91a1a}-2.91\%$
test_vmap_transformer_speed_decorator[False-True] 20.9217ms 20.4713ms 48.8490 Ops/s 50.4697 Ops/s $\color{#d91a1a}-3.21\%$
test_vmap_transformer_speed_decorator[False-False] 20.8047ms 20.4590ms 48.8782 Ops/s 50.3959 Ops/s $\color{#d91a1a}-3.01\%$
test_to_module_speed[True] 1.8557ms 1.4092ms 709.6016 Ops/s 715.4870 Ops/s $\color{#d91a1a}-0.82\%$
test_to_module_speed[False] 1.5227ms 1.3844ms 722.3582 Ops/s 729.9542 Ops/s $\color{#d91a1a}-1.04\%$
test_tc_init 0.4784ms 41.7427μs 23.9563 KOps/s 23.5649 KOps/s $\color{#35bf28}+1.66\%$
test_tc_init_tensor_only 0.1131ms 9.1193μs 109.6577 KOps/s 109.1194 KOps/s $\color{#35bf28}+0.49\%$
test_tc_init_nested 0.4970ms 82.5939μs 12.1074 KOps/s 11.9726 KOps/s $\color{#35bf28}+1.13\%$
test_tc_init_many_fields 0.1162ms 15.4560μs 64.6999 KOps/s 64.5234 KOps/s $\color{#35bf28}+0.27\%$
test_tc_first_layer_tensor 21.4000μs 1.6900μs 591.7268 KOps/s 586.2130 KOps/s $\color{#35bf28}+0.94\%$
test_tc_first_layer_tensor_only 21.0238μs 0.3807μs 2.6264 MOps/s 2.5964 MOps/s $\color{#35bf28}+1.16\%$
test_tc_first_layer_tensor_set 0.4165ms 3.6648μs 272.8633 KOps/s 270.4794 KOps/s $\color{#35bf28}+0.88\%$
test_tc_first_layer_tensor_only_set 12.7400μs 3.1417μs 318.3032 KOps/s 322.6605 KOps/s $\color{#d91a1a}-1.35\%$
test_tc_first_layer_nontensor 0.1101ms 5.8326μs 171.4514 KOps/s 173.9384 KOps/s $\color{#d91a1a}-1.43\%$
test_tc_second_layer_tensor 29.7700μs 4.1450μs 241.2570 KOps/s 242.6744 KOps/s $\color{#d91a1a}-0.58\%$
test_tc_second_layer_nontensor 0.4445ms 8.2517μs 121.1869 KOps/s 121.8601 KOps/s $\color{#d91a1a}-0.55\%$
test_unbind 0.2517s 15.6090ms 64.0656 Ops/s 70.4606 Ops/s $\textbf{\color{#d91a1a}-9.08\%}$
test_full_like 5.1210ms 4.3895ms 227.8159 Ops/s 228.1506 Ops/s $\color{#d91a1a}-0.15\%$
test_zeros_like 4.8306ms 4.3781ms 228.4083 Ops/s 228.0209 Ops/s $\color{#35bf28}+0.17\%$
test_ones_like 4.8846ms 4.3896ms 227.8136 Ops/s 228.6526 Ops/s $\color{#d91a1a}-0.37\%$
test_clone 6.9485ms 6.5354ms 153.0121 Ops/s 154.1559 Ops/s $\color{#d91a1a}-0.74\%$
test_squeeze 0.5172ms 13.1541μs 76.0220 KOps/s 75.0885 KOps/s $\color{#35bf28}+1.24\%$
test_unsqueeze 0.2422ms 0.1037ms 9.6477 KOps/s 9.5938 KOps/s $\color{#35bf28}+0.56\%$
test_split 0.6067ms 0.1731ms 5.7773 KOps/s 5.8352 KOps/s $\color{#d91a1a}-0.99\%$
test_permute 0.3237ms 0.1967ms 5.0843 KOps/s 5.1647 KOps/s $\color{#d91a1a}-1.56\%$
test_stack 51.6075ms 51.1185ms 19.5624 Ops/s 19.5691 Ops/s $\color{#d91a1a}-0.03\%$
test_cat 51.4550ms 51.2110ms 19.5270 Ops/s 23.3934 Ops/s $\textbf{\color{#d91a1a}-16.53\%}$
test_sequential_tensordict 0.3382ms 0.2138ms 4.6779 KOps/s 4.6214 KOps/s $\color{#35bf28}+1.22\%$
test_sequential_graph_module 0.5711ms 0.1175ms 8.5132 KOps/s 8.6014 KOps/s $\color{#d91a1a}-1.03\%$
test_nested_tensordict 0.5699ms 0.2800ms 3.5711 KOps/s 3.4986 KOps/s $\color{#35bf28}+2.07\%$
test_nested_graph_module 0.5879ms 0.1284ms 7.7904 KOps/s 7.9236 KOps/s $\color{#d91a1a}-1.68\%$

@vmoens vmoens mentioned this pull request Mar 11, 2026
Open
@github-actions

github-actions Bot commented Mar 11, 2026

Copy link
Copy Markdown
Contributor

PR Title Label Error

Unknown or invalid prefix [DTensor].

Current title: [DTensor] Add framework adapters for Megatron and vLLM sharding

Supported Prefixes

Your PR title must start with exactly one of these prefixes (case-insensitive):

Prefix Label Applied Example
[BugFix] or [Fix] bug [BugFix] Fix memory leak in TensorDict
[Feature] Feature [Feature] Add new storage backend
[Doc] or [Docs] documentation [Doc] Update installation guide
[Refactor] Refactor [Refactor] Clean up module imports
[CI] CI [CI] Fix workflow permissions
[Test] or [Tests] Test [Test] Add unit tests for nn module
[Compile] Compile [Compile] Fix torch.compile issue
[Performance] or [Perf] Performance [Perf] Optimize tensor operations
[Deprecation] Deprecation [Deprecation] Mark old function
[Setup] setup [Setup] Update build configuration
[Distributed] or [Dist] Distributed [Distributed] Add scatter collective
[Benchmark] or [Bench] Benchmarks [Benchmark] Add compile benchmark
[Typing] or [Type] Typing [Typing] Add type stubs
[BC-breaking] or [BC] BC-breaking [BC-breaking] Remove deprecated API
[Formatting] or [Format] Formatting [Format] Fix code style
[Quality] Quality [Quality] Improve error messages

Note: Matching is case-insensitive. Common variations (singular/plural) are supported.

@vmoens
Update
99f44ba 
[ghstack-poisoned]
@github-actions

github-actions Bot commented Mar 17, 2026

Copy link
Copy Markdown
Contributor

PR Title Label Error

Unknown or invalid prefix [DTensor].

Current title: [DTensor] Add framework adapters for Megatron and vLLM sharding

Supported Prefixes

Your PR title must start with exactly one of these prefixes (case-insensitive):

Prefix Label Applied Example
[BugFix] or [Fix] bug [BugFix] Fix memory leak in TensorDict
[Feature] Feature [Feature] Add new storage backend
[Doc] or [Docs] documentation [Doc] Update installation guide
[Refactor] Refactor [Refactor] Clean up module imports
[CI] CI [CI] Fix workflow permissions
[Test] or [Tests] Test [Test] Add unit tests for nn module
[Compile] Compile [Compile] Fix torch.compile issue
[Performance] or [Perf] Performance [Perf] Optimize tensor operations
[Deprecation] Deprecation [Deprecation] Mark old function
[Setup] setup [Setup] Update build configuration
[Distributed] or [Dist] Distributed [Distributed] Add scatter collective
[Benchmark] or [Bench] Benchmarks [Benchmark] Add compile benchmark
[Typing] or [Type] Typing [Typing] Add type stubs
[BC-breaking] or [BC] BC-breaking [BC-breaking] Remove deprecated API
[Formatting] or [Format] Formatting [Format] Fix code style
[Quality] Quality [Quality] Improve error messages

Note: Matching is case-insensitive. Common variations (singular/plural) are supported.

1 similar comment
@github-actions

github-actions Bot commented Mar 17, 2026

Copy link
Copy Markdown
Contributor

PR Title Label Error

Unknown or invalid prefix [DTensor].

Current title: [DTensor] Add framework adapters for Megatron and vLLM sharding

Supported Prefixes

Your PR title must start with exactly one of these prefixes (case-insensitive):

Prefix Label Applied Example
[BugFix] or [Fix] bug [BugFix] Fix memory leak in TensorDict
[Feature] Feature [Feature] Add new storage backend
[Doc] or [Docs] documentation [Doc] Update installation guide
[Refactor] Refactor [Refactor] Clean up module imports
[CI] CI [CI] Fix workflow permissions
[Test] or [Tests] Test [Test] Add unit tests for nn module
[Compile] Compile [Compile] Fix torch.compile issue
[Performance] or [Perf] Performance [Perf] Optimize tensor operations
[Deprecation] Deprecation [Deprecation] Mark old function
[Setup] setup [Setup] Update build configuration
[Distributed] or [Dist] Distributed [Distributed] Add scatter collective
[Benchmark] or [Bench] Benchmarks [Benchmark] Add compile benchmark
[Typing] or [Type] Typing [Typing] Add type stubs
[BC-breaking] or [BC] BC-breaking [BC-breaking] Remove deprecated API
[Formatting] or [Format] Formatting [Format] Fix code style
[Quality] Quality [Quality] Improve error messages

Note: Matching is case-insensitive. Common variations (singular/plural) are supported.

@vmoens
Update
6049f81 
[ghstack-poisoned]
@github-actions

github-actions Bot commented Mar 17, 2026

Copy link
Copy Markdown
Contributor

PR Title Label Error

Unknown or invalid prefix [DTensor].

Current title: [DTensor] Add framework adapters for Megatron and vLLM sharding

Supported Prefixes

Your PR title must start with exactly one of these prefixes (case-insensitive):

Prefix Label Applied Example
[BugFix] or [Fix] bug [BugFix] Fix memory leak in TensorDict
[Feature] Feature [Feature] Add new storage backend
[Doc] or [Docs] documentation [Doc] Update installation guide
[Refactor] Refactor [Refactor] Clean up module imports
[CI] CI [CI] Fix workflow permissions
[Test] or [Tests] Test [Test] Add unit tests for nn module
[Compile] Compile [Compile] Fix torch.compile issue
[Performance] or [Perf] Performance [Perf] Optimize tensor operations
[Deprecation] Deprecation [Deprecation] Mark old function
[Setup] setup [Setup] Update build configuration
[Distributed] or [Dist] Distributed [Distributed] Add scatter collective
[Benchmark] or [Bench] Benchmarks [Benchmark] Add compile benchmark
[Typing] or [Type] Typing [Typing] Add type stubs
[BC-breaking] or [BC] BC-breaking [BC-breaking] Remove deprecated API
[Formatting] or [Format] Formatting [Format] Fix code style
[Quality] Quality [Quality] Improve error messages

Note: Matching is case-insensitive. Common variations (singular/plural) are supported.

@vmoens
Update
7bcc0cf 
[ghstack-poisoned]
vmoens added a commit that referenced this pull request Mar 17, 2026
@vmoens
[DTensor] Add framework adapters for Megatron and vLLM sharding
4965143 
Create tensordict/dtensor_adapters/ package with adapter classes that
construct ShardingDescriptors from framework-specific parameter metadata:
- MegatronSharding: reads tensor_model_parallel, partition_dim, etc.
- VLLMSharding: reads output_dim/input_dim for column/row parallel

Each adapter provides describe(name, param) for single parameters and
describe_model(model) for all parameters. Adapters are in separate
submodules to avoid hard dependencies on the respective frameworks.

ghstack-source-id: 160b912
Made-with: Cursor
ghstack-source-id: 160b912
Pull-Request: #1651
Made-with: Cursor
@github-actions

github-actions Bot commented Mar 17, 2026

Copy link
Copy Markdown
Contributor

PR Title Label Error

Unknown or invalid prefix [DTensor].

Current title: [DTensor] Add framework adapters for Megatron and vLLM sharding

Supported Prefixes

Your PR title must start with exactly one of these prefixes (case-insensitive):

Prefix Label Applied Example
[BugFix] or [Fix] bug [BugFix] Fix memory leak in TensorDict
[Feature] Feature [Feature] Add new storage backend
[Doc] or [Docs] documentation [Doc] Update installation guide
[Refactor] Refactor [Refactor] Clean up module imports
[CI] CI [CI] Fix workflow permissions
[Test] or [Tests] Test [Test] Add unit tests for nn module
[Compile] Compile [Compile] Fix torch.compile issue
[Performance] or [Perf] Performance [Perf] Optimize tensor operations
[Deprecation] Deprecation [Deprecation] Mark old function
[Setup] setup [Setup] Update build configuration
[Distributed] or [Dist] Distributed [Distributed] Add scatter collective
[Benchmark] or [Bench] Benchmarks [Benchmark] Add compile benchmark
[Typing] or [Type] Typing [Typing] Add type stubs
[BC-breaking] or [BC] BC-breaking [BC-breaking] Remove deprecated API
[Formatting] or [Format] Formatting [Format] Fix code style
[Quality] Quality [Quality] Improve error messages

Note: Matching is case-insensitive. Common variations (singular/plural) are supported.

1 similar comment
@github-actions

github-actions Bot commented Mar 17, 2026

Copy link
Copy Markdown
Contributor

PR Title Label Error

Unknown or invalid prefix [DTensor].

Current title: [DTensor] Add framework adapters for Megatron and vLLM sharding

Supported Prefixes

Your PR title must start with exactly one of these prefixes (case-insensitive):

Prefix Label Applied Example
[BugFix] or [Fix] bug [BugFix] Fix memory leak in TensorDict
[Feature] Feature [Feature] Add new storage backend
[Doc] or [Docs] documentation [Doc] Update installation guide
[Refactor] Refactor [Refactor] Clean up module imports
[CI] CI [CI] Fix workflow permissions
[Test] or [Tests] Test [Test] Add unit tests for nn module
[Compile] Compile [Compile] Fix torch.compile issue
[Performance] or [Perf] Performance [Perf] Optimize tensor operations
[Deprecation] Deprecation [Deprecation] Mark old function
[Setup] setup [Setup] Update build configuration
[Distributed] or [Dist] Distributed [Distributed] Add scatter collective
[Benchmark] or [Bench] Benchmarks [Benchmark] Add compile benchmark
[Typing] or [Type] Typing [Typing] Add type stubs
[BC-breaking] or [BC] BC-breaking [BC-breaking] Remove deprecated API
[Formatting] or [Format] Formatting [Format] Fix code style
[Quality] Quality [Quality] Improve error messages

Note: Matching is case-insensitive. Common variations (singular/plural) are supported.

@vmoens
Update
4f0a77a 
[ghstack-poisoned]
vmoens added a commit that referenced this pull request Mar 17, 2026
@vmoens
[DTensor] Add framework adapters for Megatron and vLLM sharding
309338b 
Create tensordict/dtensor_adapters/ package with adapter classes that
construct ShardingDescriptors from framework-specific parameter metadata:
- MegatronSharding: reads tensor_model_parallel, partition_dim, etc.
- VLLMSharding: reads output_dim/input_dim for column/row parallel

Each adapter provides describe(name, param) for single parameters and
describe_model(model) for all parameters. Adapters are in separate
submodules to avoid hard dependencies on the respective frameworks.

ghstack-source-id: dc395e0
Made-with: Cursor
ghstack-source-id: dc395e0
Pull-Request: #1651
Made-with: Cursor
@github-actions

github-actions Bot commented Mar 17, 2026

Copy link
Copy Markdown
Contributor

PR Title Label Error

Unknown or invalid prefix [DTensor].

Current title: [DTensor] Add framework adapters for Megatron and vLLM sharding

Supported Prefixes

Your PR title must start with exactly one of these prefixes (case-insensitive):

Prefix Label Applied Example
[BugFix] or [Fix] bug [BugFix] Fix memory leak in TensorDict
[Feature] Feature [Feature] Add new storage backend
[Doc] or [Docs] documentation [Doc] Update installation guide
[Refactor] Refactor [Refactor] Clean up module imports
[CI] CI [CI] Fix workflow permissions
[Test] or [Tests] Test [Test] Add unit tests for nn module
[Compile] Compile [Compile] Fix torch.compile issue
[Performance] or [Perf] Performance [Perf] Optimize tensor operations
[Deprecation] Deprecation [Deprecation] Mark old function
[Setup] setup [Setup] Update build configuration
[Distributed] or [Dist] Distributed [Distributed] Add scatter collective
[Benchmark] or [Bench] Benchmarks [Benchmark] Add compile benchmark
[Typing] or [Type] Typing [Typing] Add type stubs
[BC-breaking] or [BC] BC-breaking [BC-breaking] Remove deprecated API
[Formatting] or [Format] Formatting [Format] Fix code style
[Quality] Quality [Quality] Improve error messages

Note: Matching is case-insensitive. Common variations (singular/plural) are supported.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. Test

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@vmoens