torchrec

torchrec copied to clipboard

Create a new tensor arg for stride_per_key_per_rank to facilitate torch.export

2

Summary: # Context * Currently torchrec IR serializer can't handle variable batch KJT use case. * To support VBE KJT, the `stride_per_key_per_rank` field needs to be flattened as a variable...

jd7-tr

Test include stride_per_key to KJT's flatten and unflatten

7

Differential Revision: D73051959

jd7-tr

Extend FSDP1 global clipping support for optimizers other than Shampoo

1

Summary: This diff is a followup of D73474285 and lets other dense optimizers take it the `enable_global_grad_clip` optim config. When `enalbe_global_grad_clip=True` and FSDP1 is used, it would calculate global gradient...

wz337

Update the _maybe_compute_stride_kjt logic to calculate stride based off of inverse_indices for VBE KJTs.

3

Summary: See D73051959 for context. Update the `_maybe_compute_stride_kjt` logic to calculate stride based off of `inverse_indices` for VBE KJTs. Currently, stride of VBE KJT with `stride_per_key_per_rank` is calculated as the...

jd7-tr

RT: currently, the tfra community has integerated nvidia/hkv as embedding-storage. can torchrec use the lib as embedding?

hudengjunai

Add unsharded module reference to sharded modules

5

Summary: Adding a simple unsharded module reference to sharded modules. This will be used in Dynamic Sharding by `DistributedModelParallel` to reshard an already-sharded_module. As DMP is created with only one-way...

aporialiao

Any examples?

cnglen

i want to freeze EmbeddingCollection mid training, but setting the learning rate to 0.0 does not save cuda memory and computing resources. is there a method similar to the following...

threestone965

Move util function for creating VBE KJT to the VbModelInput class

1

Summary: astitled Differential Revision: D73450618

jd7-tr

We found that the throughput of embeddingbagcollection 16p is only 1.2 times that of 8p. Are there any optimization measures?

imhandsome