[feature] Add the low level SSD APIs

[anj-s]

What does this PR do?

Adds the low level SSD APIs required for representing a tensor, reading/writing to/from memory/disk. These will be used in a future PR to implement SSD offload with FSDP.

The two main concepts here are SsdTensorHandle and SsdBuffer. The SsdTensorHandle represents a tensor that can be in memory or on disk. There are APIs to enable us to set the right metadata (memory offset, file offset, file params etc.). The SsdBuffer consists of a list of SsdTensorHandles and represents all the parameters in a module (as an example). We should be able to read all tensors to and from disk as part of the SsdBuffer.

Note: This PR is a join effort between @another-pjohnson and @anj-s .

Before submitting

• Did you have fun?
  • Make sure you had fun coding 🙃
• Did you read the contributor guideline?
• Was this discussed/approved via a Github issue? (no need for typos, doc improvements)
  • N/A
• Did you make sure to update the docs?
  • N/A
• Did you write any new necessary tests?
  • N/A
• Did you update the changelog? (if needed)
  • N/A

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

[anj-s]

Friendly ping for a review!

[another-pjohnson]

This looks good to me, although it's code that I've mostly written :-)

[min-xu-ai]

It seems that pytorch nightly version is a bit flaky with the new test?

[anj-s]

Existing tests are flaky, not the new one:
tests.nn.checkpoint.test_checkpoint_activations and tests.nn.checkpoint.test_checkpoint_activations.

[min-xu-ai]

My bad. I saw "test_offload_memory" and assumed it was a new test. I triggered a rerun. I haven't seen this one fail on CI for a long time. Is it very flaky and frequently fails?

[anj-s]

My bad. I saw "test_offload_memory" and assumed it was a new test. I triggered a rerun. I haven't seen this one fail on CI for a long time. Is it very flaky and frequently fails?

No worries! No idea, am also looking at it but couldn't figure out why it was failing now. Hasn't been flaky recently. I've seen it fail once or twice but not in the last few weeks.

[min-xu-ai]

The rerun is still failing the same way. It starts to look like something changed on CI...