Simplify and fix match_fw_and_bw_saved_for_bw_proxies implementation

[IvanYashchuk]

Base PR: #1756, keeping draft mode to prevent merges into previous PR in the stack

This change is needed to unblock @jjsjann123 for #1732.

The previous implementation uses both tensors and non-tensors saved for backward in constructing old_saved_for_backward_fw but in construction of the mirrored object from the backward trace non-tensors are ignored leading to an error that @jjsjann123 observed while working on #1732.

cc @mruberry @lantiga @ali-alshaar7

[IvanYashchuk]

Failures to fix:

FAILED thunder/tests/test_inplace_copy.py::test_prim_inplace_copy_bwd_nvfuser_cuda_thunder.dtypes.bfloat16 - AssertionError
FAILED thunder/tests/test_inplace_copy.py::test_prim_inplace_copy_bwd_nvfuser_cuda_thunder.dtypes.float16 - AssertionError
FAILED thunder/tests/test_torch_compile_executor.py::test_litgpt_fabric_for_callable - AssertionError
FAILED thunder/tests/test_torch_compile_executor.py::test_torch_compile_cat_rope_single_fusion - AssertionError
FAILED thunder/tests/test_transforms.py::test_disable_params_and_buffer_check - AssertionError
FAILED thunder/tests/test_jit_general.py::test_tom_overrides_proxy[cuda] - AssertionError
FAILED thunder/tests/test_jit_general.py::test_litgpt_variants[cuda-llama1-like] - AssertionError
FAILED thunder/tests/test_sdpaex_executor.py::test_sdpa_attn_mask[True-cuda-bf16] - AssertionError
FAILED thunder/tests/test_sdpaex_executor.py::test_sdpa_attn_mask[True-cuda-f16] - AssertionError
FAILED thunder/tests/test_sdpaex_executor.py::test_sdpa_attn_mask[False-cuda-f16] - AssertionError
FAILED thunder/tests/test_dynamo.py::test_ThunderCompilerGraphBenchmarking_LlamaMLPBenchmark - torch._dynamo.exc.BackendCompilerFailed: backend='?' raised:
AssertionError: 

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True
FAILED thunder/tests/test_sdpaex_executor.py::test_sdpa_attn_mask[False-cuda-bf16] - AssertionError
FAILED thunder/tests/test_jit_general.py::test_litgpt_variants[cuda-gpt-neox-like] - AssertionError
FAILED thunder/tests/test_jit_general.py::test_litgpt_variants[cuda-long-context-like] - AssertionError
FAILED thunder/tests/test_jit_general.py::test_litgpt_variants[cuda-llama2-like] - AssertionError
FAILED thunder/tests/test_jit_general.py::test_litgpt_variants[cuda-codellama2-like] - AssertionError
= 16 failed, 2162 passed, 228 skipped, 22 xfailed, 7 xpassed, 51635 warnings in 426.08s (0:07:06) =

[IvanYashchuk]

One more test to fix:

FAILED thunder/tests/distributed/test_fsdp.py::FSDPTest::test_rematerialize_all_gather - RuntimeError: Process 1 exited with error code 10 and exception:
Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/torch/testing/_internal/common_distributed.py", line 726, in run_test
    getattr(self, test_name)()
  File "/usr/local/lib/python3.10/dist-packages/torch/testing/_internal/common_distributed.py", line 599, in wrapper
    fn()
  File "/usr/local/lib/python3.10/dist-packages/torch/testing/_internal/common_utils.py", line 3120, in wrapper
    method(*args, **kwargs)
  File "/__w/3/s/thunder/tests/distributed/test_fsdp.py", line 136, in test_rematerialize_all_gather
    self.assertTrue(all(t in result_saved_for_bwd for t in sharded_param_names))
  File "/usr/lib/python3.10/unittest/case.py", line 687, in assertTrue
    raise self.failureException(msg)
AssertionError: False is not true