Tags from torchrec

v2025.08.25.00: refactor train_pipeline_tests for tracing (#3314)

2025-08-23T20:47:20Z

refactor train_pipeline_tests for tracing (#3314)

Summary:
Pull Request resolved: #3314

# context
* previous work in torchrec train_pipeline to refactoring the utils file
* tracing was separated from the utils
* here we also move the corresponding tests into a new file.

Reviewed By: iamzainhuda

Differential Revision: D80882439

fbshipit-source-id: 572a52f651f381b084be369314fce4bead6e853d

v2025.08.18.00

2025-08-17T20:06:43Z

ReshardingAPI Host Memory Offloading and BenchmarkReshardingHandler (#…

…3291)

Summary:
Pull Request resolved: #3291

- Implements tensor offloading to host memory inside resharding API
- Add BenchmarkReshardingHandler
- generate random plans
- calls DDP reshard API by selecting random plans
- Add reset method to train_pipelines.py

Reviewed By: aporialiao

Differential Revision: D80366926

fbshipit-source-id: a137da2f36cbacf21f0c28ae83dfc6eabba29901

v2025.08.11.00

2025-08-09T02:14:16Z

Replace int(..) with torch.sym_int(...) to make it torch.export compa…

…tible (#3270)

Summary:
Pull Request resolved: #3270

As title

Reviewed By: angelayi

Differential Revision: D79912672

fbshipit-source-id: b2586764082fdcb3665f486720a0e9ee97724d09

v2025.08.04.00: Add dtype to kt_regroup input (#3250)

2025-08-02T09:25:35Z

Add dtype to kt_regroup input (#3250)

Summary:
Pull Request resolved: #3250

Currently ir_kt_regroup always returns a float32, which is incorrect if the KTRegroupAsDict's emb_type is set

Reviewed By: malaybag

Differential Revision: D79422720

fbshipit-source-id: 2e711a7dc186b80226b2631ba94af741e05df4b7

v2025.07.28.00: fix EBC optimizer size setting for virtual table (#3239)

2025-07-26T09:58:48Z

fix EBC optimizer size setting for virtual table (#3239)

Summary:
Pull Request resolved: #3239

Add missing optimizer state processing logic similar to what EC is doing here: https://fburl.com/code/hvzzs3t4, to make sure optimizer state won't use default metadata which is the virtual table size, not the actual tensor size.

Reviewed By: EddyLXJ

Differential Revision: D78950717

fbshipit-source-id: 45eca79bbff1fe498e2707b51a6845eb603bbdfd

v2025.07.21.00: YAML config support for pipeline benchmarking (#3180)

2025-07-18T20:43:57Z

YAML config support for pipeline benchmarking (#3180)

Summary:
Pull Request resolved: #3180

Added a support for YAML file configuration of the pipeline benchmarking. This feature makes easier to reproduce complex configurations without the need to CLI arguments passing.

Example `.yaml ` file should look like:
```
RunOptions:
world_size: 2

PipelineConfig:
pipeline: "sparse"
```

Also, configs can be listed in a 'flat' way as well:
```
world_size: 2
pipeline: "sparse"
```

To run, add the `--yaml_config` flag with the `.yaml` file path. Additional flags can overwrite the `yaml` file configs as well if desired.

Reviewed By: iamzainhuda

Differential Revision: D78127340

fbshipit-source-id: 3134144fcd42761f02da81a9438546cae84b4460

v2025.07.14.00: Dynamic 2D sparse parallel (#3177)

2025-07-12T04:35:52Z

Dynamic 2D sparse parallel (#3177)

Summary:
Pull Request resolved: #3177

We add the ability to set 2D parallel configuration per module (coined as Dynamic 2D parallel). This means an EBC and EC can be sharded differently on the data parallel dimension. We can have 4 replicas per EBC shard and 2 replicas per EC shard. The previous setting requires all modules to have the same replication factor.

To do this, we introduce a lightweight dataclass that is used to provide per module configurations, allowing very granular control should it be required by the user:
```python
class DMPCollectionConfig:
module: nn.Module # this is expected to be unsharded module
plan: "ShardingPlan" # sub-tree-specific sharding plan
sharding_group_size: int
use_inter_host_allreduce: bool = False
```

The dataclass is used to provide the context for each module we are sharding. And, if configured, create separate process groups and sync logic for each of these modules.

Usage is as follows, suppose we want to use a different 2D configuration for EmbeddingCollection:
```python
# create plan for model tables over Y world size
# create plan for EmbeddingCollection tables over X world size
ec_config = DMPCollectionConfig(EmbeddingCollection, embedding_collection_plan, sharding_group_size)
model = DMPCollection(
# pass in defaults args
submodule_configs = [ec_config]
)
```

Future work includes:
- making it easier for users to create seperate sharding plans per module
- per table 2D

Reviewed By: liangbeixu

Differential Revision: D76774334

fbshipit-source-id: 27c7e0bc806d8227d784461a197cd8f1c7f6adfc

v2025.07.07.00: Fix test more generally (#3165)

2025-07-07T03:15:53Z

Fix test more generally (#3165)

Summary:
Pull Request resolved: #3165

https://www.internalfb.com/diff/D77614983
attempted to fix a test, but I still see it showing up in other tests, so this fixes it in general.

Reviewed By: huydhn

Differential Revision: D77758554

fbshipit-source-id: bd390081b68fa650f1cfd6d2a93a1fbf206aaff7

v2025.06.30.00

2025-06-29T08:22:46Z

Revert D76476676: Multisect successfully blamed "D76476676: OSS Torch…

…Rec Internal MPZCH modules" for one test failure (#3146)

Summary:
Pull Request resolved: #3146

This diff reverts D76476676
D76476676: OSS TorchRec Internal MPZCH modules by lizhouyu causes the following test failure:

Tests affected:
- [cogwheel:cogwheel_tps_basic_test#test_tps_basic_latency](https://www.internalfb.com/intern/test/562950123898458/)

Here's the Multisect link:
https://www.internalfb.com/multisect/30184635
Here are the tasks that are relevant to this breakage:
T211534727: 100+ tests, 10+ build rules failing for minimal_viable_ai

The backout may land if someone accepts it.

If this diff has been generated in error, you can Commandeer and Abandon it.

Depends on D76476676

Reviewed By: lizhouyu

Differential Revision: D77502155

fbshipit-source-id: eb990251f3276372592c30a7361579e2a3639d6c

v2025.06.23.00: minor refactoring to use list comprehension (#3125)

2025-06-21T16:46:30Z

minor refactoring to use list comprehension (#3125)

Summary:
Pull Request resolved: #3125

# context
* imported from github [#1602](#1602)
* rebased on trunk

Reviewed By: spmex

Differential Revision: D77031163

fbshipit-source-id: f65168f6ab0b74eca75b72fd60ec0ea7c762f3dc