gopikrsmscs/torch-issues huggingface dataset:
Implement device parameter in Dropout2d
Labels: feature, module: nn, triaged

### 🚀 The feature, motivation and pitch

Please provide a `device` parameter to the Dropout2d layer for creating it on a target device like so:

```python
dropout = nn.Dropout2d(
    p=0.5,
    device=torch.device('cuda'),
)
```

### Alternatives

Currently, [PyTorch's Dropout2D layer](https://pytorch.org/docs/stable/generated/torch.nn.Dropout2d.html) does not support creating it on a specific device. This means that one needs first create it on the CPU and then move it to the desired device, if any. Lacking this feature results in extra code to be written by the user:

```python
dropout = nn.Dropout2d(
    p=0.5,
)
dropout1.cuda()
```

### Additional context

Creating the layer on CPU and moving it to GPU might also be slower, when compared to creation directly on the target device, as hinted at by the performance tuning guide:

<https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html#create-tensors-directly-on-the-target-device>

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki
Simulating lower memory on GPU does not indicate simulated memory in error message
Labels: module: cuda, module: memory usage, triaged, module: CUDACachingAllocator

### 🐛 Describe the bug

I'm on an 80G A100 trying to debug an OOM folks are getting running my code on a 40G, and found the `torch.cuda.set_per_process_memory_fraction(0.5)` API which I'm using to simulate lower CUDA memory conditions. 

However, upon an OOM I get an error message such as

```
  content_understanding.utils.injected_exception.NonRetryableExceptionWrapper: (OutOfMemoryError) CUDA out of memory. Tried to allocate 64.00 MiB. GPU 0 has a total capacty of 79.15 GiB of which 38.09 GiB is free. Including non-PyTorch memory, this process has 41.06 GiB memory in use. Of the allocated memory 39.10 GiB is allocated by PyTorch, and 427.46 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
```

where the simulated reduced memory is not reflected (80G is). To repro, call the above API and allocate a > 40G tensor.

### Versions

main

cc @ptrblck
fix: Flake8-BugBear code B-026 for PyTorch
Labels: triaged, open source, release notes: distributed (fsdp), ciflow/inductor

Fixes #106571

I have fixed the B-026 error codes for Flake8 tests on the codebase. Please review and tell me anything else to do.
Thanks and excited for this first contribution to PyTorch.

Also I refer this issue which introduced [B-026](https://github.com/PyCQA/flake8-bugbear/issues/286) in `pytest-bugbear` and discuss the error code.
I have a trouble with to_symmetric
Labels: needs reproduction, module: sparse, triaged

### 🐛 Describe the bug

This is my code:
```
dataset = DGraphFin(root=path, name=dataset_name, transform=T.ToSparseTensor())

nlabels = dataset.num_classes
if dataset_name in ['DGraph']:
    nlabels = 2   

data = dataset[0]
data.adj_t =  data.adj_t.to_symmetric() 
```

I got the warning : UserWarning: Sparse CSR tensor support is in beta state.
And an error:AttributeError: 'Tensor' object has no attribute 'to_symmetric'.
I don't know why. Why someone can run successfullly, but I can't?

### Versions

Stable (2.1.0)
python 3.8 
CUDA 11.8
Windows 11

cc @alexsamardzic @nikitaved @pearu @cpuhrsch @amjames @bhosmer
_foreach_copy_ supports fast copy between cpu and cuda devices.
Labels: feature, triaged, module: mta

### 🚀 The feature, motivation and pitch

The current implementation of PyTorch's `foreach_copy` operator requires the source and target tensors to be on the same device to enter the fast route in the [code](https://github.com/pytorch/pytorch/blob/e0e15a4ac61648cc8f63f0ab102c32e8884fb5d1/aten/src/ATen/native/cuda/ForeachBinaryOpList.cu#L260). 

While I understand this limitation for other operators, I believe `copy` is a special case because copying between devices is a fairly common operation. Supporting this feature could be useful.


### Alternatives

_No response_

### Additional context

_No response_

cc @crcrpar @mcarilli
Custom Tensor Instances Do Not Work With DDP
Labels: oncall: distributed

### 🐛 Describe the bug

Existing DDP implementation throws an error for [`UninitializedParameter`](https://github.com/pytorch/pytorch/blob/e0e15a4ac61648cc8f63f0ab102c32e8884fb5d1/torch/nn/parallel/distributed.py#L765). This makes sense.

However, behind the scene, `isinstance(tensor, UninitializedParameter)` translates to `UninitializedParameter.__instancecheck__(tensor)`, which is implemented [here](https://github.com/pytorch/pytorch/blob/7bcf7da3a268b435777fe87c7794c382f444e86d/torch/nn/parameter.py#L8).

The aforementioned implementation returns `True` for custom Tensor instances (because of their `_is_param` attribute). I believe the intention of this implementation is to let custom Tensor instances to be compatible with `Parameter`. But this causes these custom Tensor instances to be instances of `UninitializedParameter`, and hence incompatible with DDP.

### Versions

NA

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu
DISABLED test_fused_int_mm_mul (__main__.TestPaternMatcher)
Labels: module: rocm, triaged, module: flaky-tests, skipped, module: inductor

Platforms: rocm

This test was disabled because it is failing in CI. See [recent examples](https://hud.pytorch.org/failure/test_fused_int_mm_mul) and the most recent trunk [workflow logs](https://github.com/pytorch/pytorch/runs/17711524027).

Over the past 72 hours, it has flakily failed in 12 workflow(s).

**Debugging instructions (after clicking on the recent samples link):**
To find relevant log snippets:
1. Click on the workflow logs linked above
2. Grep for `test_fused_int_mm_mul`


Test file path: `inductor/test_pattern_matcher.py`

cc @jeffdaily @sunway513 @jithunnair-amd @pruthvistony @ROCmSupport @dllehr-amd @jataylo @hongxiayang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler
Reenable sgd benchmark
Labels: module: inductor, module: dynamo, ciflow/inductor, ciflow/inductor-perf-test-nightly, ciflow/inductor-perf-compare

This is a cleaned up version of https://github.com/pytorch/pytorch/pull/110353--we should def land the OG PR. I wanted to launch a perf job to see the time difference in the nightly so made this one.

Job: https://github.com/pytorch/pytorch/actions/runs/6538030862


cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler
[WIP] 2D optimizer change
Labels: module: fsdp, ciflow/trunk, topic: not user facing, ciflow/periodic

This PR: 
1. This PR adds a kwarg to the flat_param util in _get_unflat_views_aligned. We use _get_unflat_views_aligned() in both runtime and optim_state_dict. For optim_state_dict, we do not need to unflatten the dtensor in tp dimension for the views. 
2. Add unit test for 2d_optim_state_dict.

Changes for optim state dict loading and unit test will be in next PR.
No speedup using semi-structured sparsity
Labels: module: sparse, triaged

When I convert linear layer weights with `to_sparse_semi_structured(module.weight)`, there doesn't appear to be any speedup when I set  the dimension of the linear layer as 4096x4096 with input size being 1x4096. Pytorch is compiled with cuSparseLT. Is there any reason why there is no speedup? In addition, the sparse alternative consumes more memory than the dense one. Any hint will be helpful!

cc @alexsamardzic @nikitaved @pearu @cpuhrsch @amjames @bhosmer
Additions to functorch_scan
Labels: open source, release notes: fx, module: inductor, module: dynamo, ciflow/inductor

This PR will contain additions to the `functorch_scan` functionality initially created with the PR https://github.com/pytorch/pytorch/pull/104442

The initial changes contain:

1. Scan being moved to higher_order_ops
2. Docstring being added for scan/scan_impl
3. A testcase to torch/testing/_internal/control_flow_opinfo_db.py

The next step would be to enable inductor lowering for scan.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler @websterbei @zou3519 @vadimkantorov @ydwu4
[AOTInductor] Add UpdateConstants for AOTInductorModel
Labels: module: inductor, ciflow/inductor

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __->__ #111326
* #111325
* #111324

Summary:
Add update_constants for AOTInductorModel

Test Plan:
Included in commit

Reviewers:

Subscribers:

Tasks:

Tags:

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @aakhundov @ColinPeppler
[AOTInductor] Add test for AOTInductorModel's interface
Labels: module: inductor, ciflow/inductor

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* #111326
* __->__ #111325
* #111324

Summary:
Add a runner for AOTInductorModel for easier test.
Add test for existing AOTInductorModel's interface (create & update constants map)

Test Plan:
Commit itself is a test.

Reviewers:

Subscribers:

Tasks:

Tags:

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @aakhundov @ColinPeppler
[AOTInductor] Rename model_runner to model_container_runner
Labels: ciflow/trunk

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* #111326
* #111325
* __->__ #111324

Summary:
We rename the model_runner to model_container_runner to prepare for
adding tests of pure model without container.

Test Plan:
commit itself is a test.

Reviewers:

Subscribers:

Tasks:

Tags:
[dynamo] Make configs more deterministic
Labels: module: bc-breaking, triaged, open source, module: inductor, module: dynamo, ciflow/inductor, release notes: dynamo, topic: bc_breaking

While not strictly necessary, it will definitely make config serialization easier and more portable, which may better support future efforts and testing/debugging.

Part of https://github.com/pytorch/pytorch/issues/111235

Split out from: https://github.com/pytorch/pytorch/pull/111298

cc @ezyang @gchanan @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler @yanboliang
Remove dynamo suported check for Windows.
Labels: triaged, open source, module: dynamo, ciflow/inductor, release notes: dynamo

We've had developers using Dynamo via torch.compile and torch._dynamo.export for months on Windows via SHARK (by just deleting these lines locally). The Windows check is moved to just the inductor backend, since it is based on Triton not supporting Windows.

Progress on #90768



cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng
[dynamo] Add LazyVariableTracker
Labels: ciflow/trunk, topic: not user facing, module: dynamo, ciflow/inductor

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* #111726
* #111725
* #111415
* #111614
* #111717
* __->__ #111306



cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng
[dynamo] annotate config with `@compile_ignored`
Labels: open source, module: inductor, module: dynamo, ciflow/inductor, release notes: dynamo

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* #111300
* #111299
* #111298
* __->__ #111303

Fixes: #111221 

cc @ezyang @voznesenskym @Chillee @yanboliang
Deprecated verbose parameter in LR schedulers
Labels: triaged, open source, topic: deprecation, release notes: optim

Fixes https://github.com/pytorch/pytorch/issues/100847

This PR follows the comment in https://github.com/pytorch/pytorch/issues/100847#issuecomment-1546247239 by deprecating the `verbose` parameter and removing the print statements. Removing the print statements is technically BC breaking, so I would be okay with putting them back in.

To be less annoying, this PR raises a warning only when `verbose` is explicitly passed in.
[dynamo] guarded config
Labels: open source, module: dynamo, ciflow/inductor, release notes: dynamo

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* #111300
* __->__ #111299
* #111298
* #111303

---

Fixes: https://github.com/pytorch/pytorch/issues/110682

Replaces: https://github.com/pytorch/pytorch/pull/111074

The guards are installed based on config that is valid at the call to `torch.compile`, rather than at any subsequent call / triggered compilation. Subsequent compilations will restore the config if there is a config mismatch of the existing global config with the saved config.

TODO:
- [X] add tests

Follow up PRs:
- [x] add revised cache size computation (follow up PR: #111300 , based on: https://github.com/pytorch/pytorch/pull/107496) 
- [ ] handle run-only mode?
- [ ] config restoration itself is not thread-safe (tracked: https://github.com/pytorch/pytorch/issues/111150)

cc @voznesenskym @ezyang @Chillee
[dynamo] md5 hash non `compile_ignored` configs
Labels: open source, topic: not user facing, module: dynamo, ciflow/inductor

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* #111300
* #111299
* __->__ #111298
* #111303

fixes: https://github.com/pytorch/pytorch/issues/111235

cc @voznesenskym @ezyang @Chillee
[Quant] [PT2] Add ConvBNAdd(ReLU) Annotation into X86InductorQuantizer
Labels: open source, ciflow/trunk, release notes: quantization, module: inductor

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __->__ #111281
* #111280

**Summary**
This PR adds ConvBNAdd(ReLU) QAT Annotation into `X86InductorQuantizer`.

**Test Plan**
```
python -m pytest test_x86inductor_quantizer.py -k test_qat_conv2d_binary_with_quantizer_api
python -m pytest test_x86inductor_quantizer.py -k test_qat_conv2d_binary_unary_with_quantizer_api
python -m pytest test_mkldnn_pattern_matcher.py -k test_qat_qconv2d_add
python -m pytest test_mkldnn_pattern_matcher.py -k test_qat_qconv2d_add_relu
```


cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler
[Quant] [PT2] Enable QAT Quantization flow in X86InductorQuantizer
Labels: open source, ciflow/trunk, release notes: quantization, module: inductor

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* #111281
* __->__ #111280

**Summary**
This PR enables PT2 QAT Quantization flow in `X86InductorQuantizer`.



**Test Plan**
```
python -m pytest test_x86inductor_quantizer.py -k test_qat_conv2d_with_quantizer_api
python -m pytest test_x86inductor_quantizer.py -k test_qat_conv2d_unary_with_quantizer_api
python -m pytest test_mkldnn_pattern_matcher.py -k test_qat_qconv2d
python -m pytest test_mkldnn_pattern_matcher.py -k test_qat_qconv2d_relu
```

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler
Specialize on symint if used as a dict key
Labels: module: dynamo, ciflow/inductor

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __->__ #111263

[WIP] Needs tests - specialize on symint if used as a dict key

cc @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng
[AOTInductor] 14k models: AssertionError: Dynamo attempts to add additional input during export
Labels: triaged, oncall: pt2

### 🐛 Describe the bug

53 errors like: AssertionError: Dynamo attempts to add additional input during export: value=0.1767766952966369, source=NNModuleSource(base=AttrSource(base=NNModuleSource(base=AttrSource(base=LocalSource(local_name='self', cell_or_freevar=False), member='wscale')), member='scale')) (example ./generated/test_taesungp_contrastive_unpaired_translation.py:Generator # pytest ./generated/test_taesungp_contrastive_unpaired_translation.py -k test_007)

repro:

```
git clone https://github.com/jansel/pytorch-jit-paritybench.git
# set PYTHONPATH like the following example
PYTHONPATH=/home/yangche/work/pytorch-jit-paritybench/paritybench:$PYTHONPATH python main.py --compile_mode=aot_inductor -e ./generated/test_taesungp_contrastive_unpaired_translation.py
```

### Versions

Collecting environment information...
PyTorch version: 2.2.0a0+git9e8da56
Is debug build: False
CUDA used to build PyTorch: 12.0
ROCM used to build PyTorch: N/A

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305
[AOTInductor] 14k models: AssertionError: Failed to produce a graph during tracing. Tracing through 'f' must produce a single graph.
Labels: triaged, oncall: pt2

### 🐛 Describe the bug

169 errors like: AssertionError: Failed to produce a graph during tracing. Tracing through 'f' must produce a single graph. (example ./generated/test_alibaba_FederatedScope.py:DummyRegularizer # pytest ./generated/test_alibaba_FederatedScope.py -k test_002)

repro:

```
git clone https://github.com/jansel/pytorch-jit-paritybench.git
# set PYTHONPATH like the following example
PYTHONPATH=/home/yangche/work/pytorch-jit-paritybench/paritybench:$PYTHONPATH python main.py --compile_mode=aot_inductor -e ./generated/test_alibaba_FederatedScope.py
```

### Versions

Versions
Versions
Collecting environment information...
PyTorch version: 2.2.0a0+git9e8da56
Is debug build: False
CUDA used to build PyTorch: 12.0
ROCM used to build PyTorch: N/A

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305
[AOTInductor] 14k models: UserError: Tried to use data-dependent value in the subsequent computation
Labels: triaged, oncall: pt2

### 🐛 Describe the bug

40 errors like: UserError: Tried to use data-dependent value in the subsequent computation. This can happen when we encounter unbounded dynamic value that is unknown during tracing time.You will need to explicitly give hint to the compiler. Please take a look at constrain_as_value OR constrain_as_size APIs (example ./generated/test_wasidennis_AdaptSegNet.py:CrossEntropy2d # pytest ./generated/test_wasidennis_AdaptSegNet.py -k test_002)

How to repro:

```
git clone https://github.com/jansel/pytorch-jit-paritybench.git
# set PYTHONPATH like the following example
PYTHONPATH=/home/yangche/work/pytorch-jit-paritybench/paritybench:$PYTHONPATH python main.py --compile_mode=aot_inductor -e ./generated/test_wasidennis_AdaptSegNet.py
```



### Versions

Versions
Versions
Collecting environment information...
PyTorch version: 2.2.0a0+git9e8da56
Is debug build: False
CUDA used to build PyTorch: 12.0
ROCM used to build PyTorch: N/A

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305
[AOTInductor] 14k models: AssertionError: original output #2 is None, but only the following types are supported
Labels: triaged, oncall: pt2

### 🐛 Describe the bug

85 errors like: AssertionError: original output # 1 is None, but only the following types are supported: (<class 'torch.Tensor'>, <class 'torch.SymInt'>, <class 'torch.SymFloat'>, <class 'torch.SymBool'>) (example ./generated/test_nlp_uoregon_trankit.py:SelfAttention # pytest ./generated/test_nlp_uoregon_trankit.py -k test_031)

How to repro:

```
git clone https://github.com/jansel/pytorch-jit-paritybench.git
# set PYTHONPATH like the following example
PYTHONPATH=/home/yangche/work/pytorch-jit-paritybench/paritybench:$PYTHONPATH python main.py --compile_mode=aot_inductor -e ./generated/test_nlp_uoregon_trankit.py
```

### Versions

Versions
Collecting environment information...
PyTorch version: 2.2.0a0+git9e8da56
Is debug build: False
CUDA used to build PyTorch: 12.0
ROCM used to build PyTorch: N/A

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305
[dynamo] `ConfigModule`: Implement mechanism to hash non-`compile_ignored` configs quickly
Labels: triaged, oncall: pt2

### 🚀 The feature, motivation and pitch

Tracked: https://github.com/pytorch/pytorch/issues/111220

compile_ignored defined by https://github.com/pytorch/pytorch/issues/111221
Hash lazily using an `_is_dirty` flag, use sorted dict `_repr_` str as hash input, using MD5, output is a binary string.

`_is_dirty` is only flagged for non-`compile_ignored` keys

There should also be a thread-local variant as per: https://github.com/pytorch/pytorch/issues/111150

Additional Work:
- (non-blocking) Make the hash deterministic


cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305
torch.export: cannot instantiate Dim from REPL
Labels: fb-exported, module: inductor, ciflow/inductor

Summary:
```
In [1]: import torch
   ...: torch.export.Dim('foo', min=1, max=16)
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[1], line 2
      1 import torch
----> 2 torch.export.Dim('foo', min=1, max=16)

File /..../torch/export/__init__.py:319, in Dim(name, min, max)
    317 assert _max > _min, f"Cannot create Dim with inconsistent min={min}, max={max}"
    318 dim = _Dim(name, (int,), {"min": _min, "max": _max})
--> 319 dim.__module__ = inspect.getmodule(inspect.stack()[1][0]).__name__  # type: ignore[union-attr]
    320 return dim

AttributeError: 'NoneType' object has no attribute '__name__'
```

Test Plan: Repeat above repro

Differential Revision: D50275165



cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler
[TEST] Test is_allowed at the end
Labels: module: dynamo, ciflow/inductor

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __->__ #111228

Signed-off-by: Edward Z. Yang <ezyang@meta.com>

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng
nanogpt_generate: C++ compile times out, because the generated .cpp file is too large.
Labels: triaged, oncall: pt2

`nanogpt_generate` exists in torchbench before https://github.com/pytorch/benchmark/pull/1911,

Repro:
```
benchmarks/dynamo/torchbench.py --bfloat16 --accuracy --inference --device cuda --export-aot-inductor  --only nanogpt_generate
[AOTInductor] 14K models: TypeError: make_boxed_func..g() missing 1 required positional argument: 'args'
Labels: triaged, oncall: pt2

### 🐛 Describe the bug

386 errors like: TypeError: make_boxed_func..g() missing 1 required positional argument: 'args' (example ./generated/test_Yang_Liu1082_InvDN.py:_UpsampleBlock # pytest ./generated/test_Yang_Liu1082_InvDN.py -k test_014)

How to repro:

```
git clone https://github.com/jansel/pytorch-jit-paritybench.git
# set PYTHONPATH like the following example
PYTHONPATH=/home/yangche/work/pytorch-jit-paritybench/paritybench:$PYTHONPATH python main.py --compile_mode=aot_inductor -e ./generated/test_pyro_ppl_pyro.py 
```

### Versions

Collecting environment information...                                                                                                                             
PyTorch version: 2.2.0a0+git9e8da56                                                                                                                               
Is debug build: False                                                                                                                                             
CUDA used to build PyTorch: 12.0                                                                                                                                  
ROCM used to build PyTorch: N/A

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305
[dynamo] Investigate interop issues with torch_scatter/torch_sparse/pyg_lib
Labels: oncall: pt2, module: dynamic shapes, module: dynamo

In [this post](https://pytorch-geometric.readthedocs.io/en/latest/advanced/compile.html) pytorch-geometric mentions they needed to disable `torch_scatter`, `torch_sparse`, and `pyg_lib` to get `torch.compile` working.

We should investigate what is causing these libraries to fail, create issues for the underling problems, and fix them.

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng
[dynamo] annotate configs with `@compile_ignored` 
Labels: triaged, module: dynamo

### 🚀 The feature, motivation and pitch

Originated: https://github.com/pytorch/pytorch/pull/111074#discussion_r1358455363

Tracked: https://github.com/pytorch/pytorch/issues/111220

This annotation means that we will not recompile if these configs change

We will also store separate dicts instead of a single `_config`:
`_compile_config` and `_compile_ignored`.

We can thus more quickly read and hash this config.

---

Currently ignorable dynamo configs:
```python
dynamo_guarded_config_ignorelist = {
    "log_file_name",
    "verbose",
    "verify_correctness",  # will not affect model, will raise RuntimeError
    # (no silent change to compilation behaviour)
    "cache_size_limit",
    "accumulated_cache_size_limit",
    "print_specializations",
    "replay_record_enabled",
    "cprofile",  # only wraps _compile, not graph
    "repro_after",
    "repro_level",
    "repro_forward_only",
    "repro_tolerance",
    "same_two_models_use_fp64",
    "error_on_recompile",  # safe because: will throw error
    "report_guard_failures",
    "report_all_guard_failures",
    "base_dir",  # used for minifying / logging
    "translation_validation",
    "translation_validation_timeout",
    "translation_validation_no_bisect",
    "DEBUG_DIR_VAR_NAME",
    "debug_dir_root",
}
```


cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng
[dynamo] Tracking: improve `ConfigModule`
Labels: triaged, oncall: pt2

### Tracked
- [ ] https://github.com/pytorch/pytorch/issues/111150
- [ ] https://github.com/pytorch/pytorch/issues/111221
- [ ] https://github.com/pytorch/pytorch/issues/111235

### 🚀 The feature, motivation and pitch



Part of making config more portable between compilations: https://github.com/pytorch/pytorch/issues/110682

Affects consumers:
- dynamo
- inductor
- functorch



cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305
[dynamo] annotate `allow_in_graph` with soft constraints
Labels: module: optimizer, triaged, oncall: pt2, module: dynamo

### 🚀 The feature, motivation and pitch

Originally suggested: https://github.com/pytorch/pytorch/pull/110709/files#r1358391537

```python
@allow_in_graph(
    constraints={
        "mutates": "only self.state, and where Tensors are concerned, only add new (not remove/mutate)",
        "side-effects": "no side-effects like printing or modifying global variables",
        "return values": "must be derived from self.params (guarded) tensor metadata",
    }
)
def _init_group(self, ..):
  ..
```

The constraints cannot be enforced of course - dynamo cannot inspect the interior of the annotated function. But we are explicitly asking the user to guarantee these properties hold. 

It is still quite unsafe, but at least we know why it is unsafe, and the user might understand why something might break if they try to incorrectly modify this `@allow_in_graph` function.

---

This is stronger than the implicit contract suggested in: https://github.com/pytorch/pytorch/blob/35750bf9d134726b47451d78b731017ad27fbc8b/docs/source/torch.compiler_faq.rst#L497
> By using allow_in_graph to annotate a function, you must make sure your code meets the following requirements:
>
> - All outputs in your function only depend on the inputs and do not depend on any captured Tensors.
> - Your function is functional. That is, it does not mutate any state. This may be relaxed; we actually support functions that appear > to be functional from the outside: they may have in-place PyTorch operations, but may not mutate global state or inputs to the function.
> - Your function does not raise data-dependent errors.

```diff
- All outputs in your function only depend on the inputs and do not depend on any captured Tensors.
+ All outputs in your function only depend on the inputs or Tensor metadata and do not depend on any captured Tensor values.
```

cc @vincentqb @jbschlosser @albanD @janeyx99 @crcrpar @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @zou3519 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng
[WIP] Support tensors as Dict keys
Labels: triaged, open source, module: dynamo, ciflow/inductor, release notes: dynamo

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* #108420
* #110524
* __->__ #111196
* #110523
* #110522

This prepares the PR where we implement sets in terms of dicts.

While implementing this I had to make hashable keys somewhat consistent
all throughout. I did that by having an auxiliary class (similar to the
one that was inside SetVariable). This variable is opaque on purpose, so
that it fails hard if it unadvertedly leaks back into user code.

We also found and fixed a number of latent bugs, like where dict with
builtin types was going through to UserDefinedObjectVariable, even
though we had a test that was testing that this worked (the test was not
testing whether we were tracing anything or not).

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng
Trace dynamic batch size with make_fx
Labels: actionable, module: fx, oncall: pt2, oncall: fx, module: functorch, module: dynamic shapes

Hello,

is there a way to keep the batch dimension dynamic when tracing with `make_fx` as described in the README.md.

Example

```
import torch
from functorch import make_fx
from torch.func import vmap, jacrev

model = torch.nn.Sequential(
    torch.nn.Linear(2, 512),
    torch.nn.Tanh(),
    torch.nn.Linear(512, 2),
)

trace_inp = torch.randn(1, 2)
traced_model = make_fx(vmap(jacrev(model)), tracing_mode='real')(trace_inp)

# This will work
print(traced_model(torch.randn(1, 2)))

# This will fail with: RuntimeError: shape '[1, 2]' is invalid for input of size 4
print(traced_model(torch.randn(2, 2)))

```

Switching to `tracing_mode='symbolic'` will result in `RuntimeError: Cannot call sizes() on tensor with symbolic sizes/strides`

Appreciate the help!

Best
Tim

cc @ezyang @SherlockNoMad @EikanWang @jgong5 @wenzhe-nrv @msaroufim @wconstab @bdhirsh @anijain2305 @zou3519 @Chillee @samdow @kshitij12345 @janeyx99
Fix flake8-bugbear B019
Labels: module: cpu, triaged, open source, release notes: fx, topic: not user facing, ciflow/mps, module: inductor, ciflow/inductor, module: export

First time contributor helping out with issue #106571


cc @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10 @voznesenskym @penguinwu @EikanWang @Guobing-Chen @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan
Guards elimination for unused variables
Labels: triaged, oncall: pt2

### 🚀 The feature, motivation and pitch

For example 
`def f(m : Module, a : Tensor):
       return m.para + a
`
when this function compiled by dynamo, with config.guard_nn_modules true, totally three guards for m, m.para and a respectively will be generated. 

In fact, m itself is not used directly which incurs a NN_MODULE guard. If this function is called with different modules and the para field within them are of the same shape, perhaps the guard for m could be eliminated for less check and better hit chance.

Consider two implementation approaches:
1. Modify bytecode handlers in InstructionTranslatorBase and call_function/method in different VariableTracker subclasses to do some liveness analysis for variable guard and do elimination.
2. Perform analysis directly on compiled fx graphs and collect live variable guards, then use this collection as the base to generate GuardedCode.

Any suggestions? Thanks!

### Alternatives

_No response_

### Additional context

_No response_

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305
[ONNX][Exporter] Maintain support for exporter arguments export_params and keep_initializers_as_inputs
Labels: module: onnx, triaged

### 🚀 The feature, motivation and pitch

The [`ExportOptions`](https://pytorch.org/docs/main/onnx_dynamo.html#torch.onnx.ExportOptions) supported by the dynamo based exporter does not support the arguments `export_params` and `keep_initializers_as_inputs` that were previously supported in the torch-script based exporter.

These arguments are useful for scenarios where the data associated with the parameters in the exported model are actually sourced from elsewhere (for example from the `torch.nn.Module` in the case of [`ORTModule`](https://github.com/microsoft/onnxruntime/blob/dad70ad4e80fb6461925d33f331d26ff9c52bb6c/orttraining/orttraining/python/training/ortmodule/ortmodule.py#L27).

The lack of support for these arguments implies that we end up with a duplicated copy of the model parameters and need to manually move remove them in the exported model. This creates a memory bloat. The problem is further magnified if we have a distributed training scenario where the model is exported for each of the devices that share the same host (where we have multiple copies of the same exported model). This can be a bottleneck for large models.

If the exporter supports these arguments, the exported model will be considerably smaller in size and will not lead to a bloat in the host memory.

cc @BowenBao 

### Alternatives

_No response_

### Additional context

_No response_
Reshape mask for triu to have the same shape as input.
Labels: triaged, open source

Currently this gives this warning:
[W Resize.cpp:35] Warning: An output with one or more elements was resized since it had shape [10, 10], which does not match the required output shape [1, 12, 10, 10]. This behavior is
deprecated, and in a future PyTorch release outputs will not be resized unless they have zero elements. You can explicitly reuse an out tensor t by resizing it, inplace, to zero elements with t.resize_(0). (function _resize_output_check)

Fixes #ISSUE_NUMBER
[symnode] support symbool -> symint casting for symint arithmetic
Labels: open source, release notes: fx, module: inductor, ciflow/inductor

Possible hacky fix to: https://github.com/pytorch/pytorch/issues/110738

I doubt you would approve @lezcano, but sympy does not handle bool -> int casting as python does.

I don't think this is complete, but it is at least a start so I'll leave it up here in case anyone has a better idea for how to continue. 


cc @lezcano @peterbell10
add test for consecutive aot inductor compiles
Labels: fb-exported, topic: not user facing, module: inductor

Differential Revision: D50246956




cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler
Result of adding noise is very different in mps vs cuda or cpu
Labels: triaged, module: mps

### 🐛 Describe the bug

Possibly related to #84936 I was trying to follow [this tutorial](https://colab.research.google.com/drive/1NFxjNI-UIR7Ku0KERmv7Yb_586vHQW43?usp=sharing#scrollTo=F724QP16w7-A), but I noticed the results of adding noise in `cpu/cuda` are quite different than `mps`. This was tested on both `PyTorch 2.2.0.dev20231008` and `PyTorch 2.1.0`.

![cpu_mps](https://github.com/pytorch/pytorch/assets/5020822/da099f7e-d353-4fea-ba47-340c7080c703)

I created a gist to easily replicate the problem here:

https://gist.github.com/atabakd/e4fa20c23f7a6059dac794a82dccfb38

### Versions


```
Collecting environment information...
PyTorch version: 2.1.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.0 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.0.40.1)
CMake version: version 3.27.4
Libc version: N/A

Python version: 3.9.18 | packaged by conda-forge | (main, Aug 30 2023, 03:53:08)  [Clang 15.0.7 ] (64-bit runtime)
Python platform: macOS-14.0-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M1 Max

Versions of relevant libraries:
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.24.4
[pip3] torch==2.1.0
[pip3] torchaudio==2.1.0
[pip3] torchvision==0.16.0
[conda] numpy                     1.24.4           py39h485cf63_0    conda-forge
[conda] pytorch                   2.1.0                   py3.9_0    pytorch
[conda] torchaudio                2.1.0                  py39_cpu    pytorch
[conda] torchvision               0.16.0                 py39_cpu    pytorch
```


cc @kulinseth @albanD @malfet @DenisVieriu97 @razarmehr @abhudev
Regression on CUDA 12.1 for vanilla transformer layer
Labels: high priority, triage review, module: performance, module: cuda

### 🐛 Describe the bug

Hi, we are benchmarking [deep vit](https://github.com/lucidrains/vit-pytorch/blob/main/vit_pytorch/deepvit.py) models and found out the CUDA 12.1 binary are actually slower than CUDA 11.8 for PyTorch 2.1.0 release. This is not expected as we see better perf on other LLMs like Megatron(GPT) and OPT. One potential reason is those repos are using flash-attention or xformers which benefit from CUDA 12 and transformer engine on H100 GPUs. While deep VIT is using vanilla implementation: https://github.com/lucidrains/vit-pytorch/blob/main/vit_pytorch/deepvit.py. But we are not expecting a regression. Any ideas? thanks!

here is the output showing ~10% regression on single node
PT 2.1.0 + CUDA 11.8
```
[1,0]<stdout>:time taken for the last 1 steps is 2.420312336000279, the throughput is 13.221434078579314 images/stime taken for forward pass is 0.261821381000118time taken for backward pass is 1.758593597000072
[1,0]<stdout>:Average elapsed time is 2.3975933989263996, throughput is 13.34671675953439peak allocated memory is 25.845419008Average forward time is 0.24408430756846333Average backward time is 1.75420702246313
```

PT 2.1.0 + CUDA 12.1

```
[1,0]<stdout>:time taken for the last 1 steps is 2.7207883900000525, the throughput is 11.76129687910032 images/stime taken for forward pass is 0.26708757999995214time taken for backward pass is 1.9990846610003246
[1,0]<stdout>:Average elapsed time is 2.72146025645262, throughput is 11.75839328321167peak allocated memory is 25.845419008Average forward time is 0.26492700765261123Average backward time is 2.001053710136823

```

attached training script and reproduciable steps

```
pip install vit-pytorch
pip install packaging
torchrun --nnodes=1 --nproc_per_node=8 train.py
```



### Versions

PyTorch 2.1.0
cuda 11.8 vs cuda 12.1
tested on AWS P5.48xlarge

cc @ezyang @gchanan @zou3519 @kadeng @ptrblck
[dtensor] fix dtype/device conversion on nn.Modules
Labels: ciflow/trunk, release notes: distributed (dtensor)

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __->__ #111162

This PR fixes the dtype conversion issue when calling nn.Module.to after
swapping parameters with DTensors, the problem was nn.Module._apply by
default uses a `.data` setting way to swap out storages directly instead
of just replacing the nn.Parameters. For wrapper tensor subclasses, it
would not work as wrapper tensor subclass is different from normal
tensor as its real storage points separately, i.e.

* DTensor wrapper subclass instance `storage` does not point to real storage
* The real storage of DTensor pointed to its local shard. But if we do `.data=` it seems that the change will not propagate into the DTensor's local shard

This probably also fix https://github.com/pytorch/pytorch/issues/102812
overloads can perhaps be more performant?
Labels: module: optimizer, triaged, module: dispatch, module: codegen

After I land the 0dim overload in #111079, foreach_add will have the following signatures, generated in the following order:
```
 static PythonArgParser parser({
    "_foreach_add_(TensorList self, Scalar scalar)",
    "_foreach_add_(TensorList self, ScalarList scalars)",
    "_foreach_add_(TensorList self, Tensor other, Scalar alpha=1)",
    "_foreach_add_(TensorList self, TensorList other, *, Scalar alpha=1)",
  }, /*traceable=*/false);
```

The order is probably due to Scalars being more performant on CUDA.

## Problem:
The whole reason I added a tensor overload is so that in the CPU path, we don't continuously rewrap tensors as we dispatch into add(...). However, the above ordering means that code like `_foreach_add_(TensorList tensors, Tensor scalar_tensor)`, where 
`scalar_tensor` could be `torch.tensor(1.0)`, the first overload will "suffice", `scalar_tensor` will get casted to a normal Scalar, and we fall back into our loop of slowness.

Instead, I'd prefer `scalar_tensor` to not get coerced at all, and that the overload that ends up getting called should be `_foreach_add_(TensorList self, Tensor other, Scalar alpha=1)`. 

## Proposal:
I wonder if there's a way to tweak the parser code to penalize for coercion. It should check all the overloads and pick the one with the smallest coercion necessary before dispatching to it. Maybe the key here is to weigh `overloaded_args`. https://github.com/pytorch/pytorch/blob/6748a14a71f235501cd050e15d1810f7a0a75dbd/torch/csrc/utils/python_arg_parser.cpp#L1655-L1660

Con: I was just made aware that this can be detrimental to dispatcher performance due to slow is_instance checks...so maybe this would not be more performant after all...

Workaround:
My workaround today is just explicitly using an alpha to force the overload to be called :)

cc @vincentqb @jbschlosser @albanD @crcrpar @ezyang @bhosmer @bdhirsh @kadeng
[wip] Add wanda pruner to torch.ao.pruning
Labels: release notes: AO frontend

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __->__ #111156

Summary:

This PR adds in support for Wanda pruning into pytorch.
Many thanks to @Eric-mingjie who provided
[inital design and POC](https://github.com/Eric-mingjie/pytorch/tree/patch-1).

[Wanda](https://arxiv.org/abs/2306.11695) is a pruning method that takes into
acccount both the weight and activation statistics when pruning.

In order to keep track of activation norm statistics, this sparsifier
utilizes the `torch.ao.quantization` API and a custom observer,
`PerChannelNormObserver`.

A user can then calibrate their model by doing
```
for x in dataset:
    model(x)
```

We only use the observer to colllect input norm statistics, which we
access when calling `update_mask`

Finally, we remove observers and mask parameterization by calling
`squash_mask()`.

Note that at this time we only support using the unpruned activations
for calibration.

WandaSparsifier supports both unstructured and semi-structured (2:4)
pruning modes.

Test Plan:
```
pytest test/ao/sparsity/test_wanda_sparsifier.py
```

Reviewers:

Subscribers:

Tasks:

Tags:
[HigherOrderOp] Move _map.py to _higher_order_ops
Labels: ciflow/trunk, topic: not user facing, module: export

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __->__ #111152
* #111147
* #111092



cc @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan

Differential Revision: [D50332159](https://our.internmc.facebook.com/intern/diff/D50332159)
[dynamo] `ConfigModule` and `config.patch` are not thread safe
Labels: triaged, oncall: pt2

### 🐛 Describe the bug

Config is a global object that can be accessed concurrently by multiple threads. Hence
1. patching should be performed on thread-local values
2. (optional) modifying/reading the global copy should be based on RWLocks.

When accessing the config object, if a thread-local exists (it is patched) then use that. Else, use the global.

### Priority
low now, higher with https://github.com/pytorch/pytorch/pull/111074, when we expect a lot more patching due to `torch.compiled` callables being able to restore a previously saved config whenever you call them.


### Versions

main, https://github.com/pytorch/pytorch/pull/111074

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305
[HigherOrderOp] make MapHigherOrderOp error out when graph break
Labels: ciflow/trunk, module: dynamo, ciflow/inductor, release notes: dynamo

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* #111152
* __->__ #111147
* #111092



cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng

Differential Revision: [](https://our.internmc.facebook.com/intern/diff/)
[pytorch][PR] [Inductor][FX passes] Pre grad batch relu fusion
Labels: fb-exported, module: inductor, ciflow/inductor

Summary: We detect independent relu operators and do the fusion in the pre grad.

Test Plan:
### unit test
```
buck2 test mode/dev-nosan //caffe2/test/inductor:group_batch_fusion
```
Test UI: https://www.internalfb.com/intern/testinfra/testrun/16888498608558485

### Inlinve cvr
f479655232
```
buck2 run mode/opt //scripts/jackiexu0313/pt2:local_model_with_pt2 -- --test_mode split_batch_group
```
before vs after transformation
https://www.internalfb.com/intern/diffing/?paste_number=851907099

```
buck2 run mode/opt //scripts/jackiexu0313/pt2:local_model_with_pt2 -- --test_mode split_batch_group -c
```

P852036786

Differential Revision: D50207610




cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler
[dynamo/higher order op] fix flaky / disabled tests - context fn is not `None` when a `noop_context_fn`
Labels: triaged, open source, topic: not user facing, module: dynamo, ciflow/inductor

test/dynamo/test_activation_checkpointing.py fails without this

Fails on main with
`_experimental_support_context_fn_in_torch_utils_checkpoint = True`


cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng @ydwu4