[Performance][Optimizer] Enable using UVA and FP16 with SparseAdamOptimizer

[nv-dlasalle]

Description

This PR enables transferring optimizer states via UVA (by default) as well as storing them in FP16 (requires opt-in). While there are many factors, combining both of these optimizations improves performance in backward by about 2x and significantly cuts down on memory usage (both from FP16, and from not need allocate buffer tensors to copy to from the GPU).

This depends on #3997 to ensure the UVA arrays get properly freed when the optimizer is destroyed.

Checklist

Please feel free to remove inapplicable items for your PR.

• The PR title starts with [$CATEGORY] (such as [NN], [Model], [Doc], [Feature]])
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented
• To the best of my knowledge, examples are either not affected by this change,
  or have been fixed to be compatible with this change

Changes

• Adds "Scatter" series of UVA function for scatter data from the GPU into pinned CPU memory.
• Adds parameters use_uva and dtype to the SparseAdamOptimizer.

[dgl-bot]

To trigger regression tests:

• @dgl-bot run [instance-type] [which tests] [compare-with-branch];
  For example: @dgl-bot run g4dn.4xlarge all dmlc/master or @dgl-bot run c5.9xlarge kernel,api dmlc/master

[classicsong]

LTGM

[jermainewang]

The test_unified_tensor UT failed:

tests/pytorch/test_unified_tensor.py::test_unified_tensor FAILED         [ 99%]
tests/pytorch/test_unified_tensor.py::test_multi_gpu_unified_tensor[1] PASSED [ 99%]
tests/pytorch/test_unified_tensor.py::test_multi_gpu_unified_tensor[2] SKIPPED [100%]

=================================== FAILURES ===================================
_____________________________ test_unified_tensor ______________________________

    @unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
    @unittest.skipIf(F.ctx().type == 'cpu', reason='gpu only test')
    def test_unified_tensor():
        test_row_size = 65536
        test_col_size = 128
    
        rand_test_size = 8192
    
        input = th.rand((test_row_size, test_col_size))
        input_unified = dgl.contrib.UnifiedTensor(input, device=th.device('cuda'))
    
        seq_idx = th.arange(0, test_row_size)
        assert th.all(th.eq(input[seq_idx], input_unified[seq_idx]))
    
        seq_idx = seq_idx.to(th.device('cuda'))
        assert th.all(th.eq(input[seq_idx].to(th.device('cuda')), input_unified[seq_idx]))
    
        rand_idx = th.randint(0, test_row_size, (rand_test_size,))
        assert th.all(th.eq(input[rand_idx], input_unified[rand_idx]))
    
>       rand_idx = rand_idx.to(th.device('cuda'))
E       RuntimeError: CUDA error: invalid argument
E       CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
E       For debugging consider passing CUDA_LAUNCH_BLOCKING=1.

tests/pytorch/test_unified_tensor.py:36: RuntimeError

I saw you've changed the IndexSelect operator so probably related.

[yaox12]

@nv-dlasalle If you don't have anything to add, I think it's OK to merge this PR.

[dgl-bot]

Commit ID: b16bb40

Build ID: 22

Status: ✅ CI test succeeded

Report path: link

Full logs path: link