DeepSpeed

1
# Copyright (c) Microsoft Corporation.
2
# SPDX-License-Identifier: Apache-2.0
3

4
# DeepSpeed Team
5

6
import torch
7
from deepspeed.pt.deepspeed_linear import LinearModuleForZeroStage3
8
from deepspeed.pt.log_utils import logger
9
from deepspeed.accelerator import get_accelerator
10

11

12
def see_memory_usage(message):
13

14
    # Print message except when distributed but not rank 0
15
    logger.info(message)
16
    logger.info(
17
        "Memory Allocated %s GigaBytes ",
18
        get_accelerator().memory_allocated() / (1024 * 1024 * 1024),
19
    )
20
    logger.info(
21
        "Max Memory Allocated %s GigaBytes",
22
        get_accelerator().max_memory_allocated() / (1024 * 1024 * 1024),
23
    )
24
    logger.info(
25
        "Cache Allocated %s GigaBytes",
26
        get_accelerator().memory_cached() / (1024 * 1024 * 1024),
27
    )
28
    logger.info(
29
        "Max cache Allocated %s GigaBytes",
30
        get_accelerator().max_memory_cached() / (1024 * 1024 * 1024),
31
    )
32

33

34
tens = torch.rand(1024, 16384, dtype=torch.half, device=torch.device(get_accelerator().device_name()))
35
tens_back = tens.detach().clone()
36

37
#linear_bk = torch.nn.functional.linear
38
#torch.nn.functional.linear = deepspeed.pt.deepspeed_linear.LinearFunctionForZeroStage3.apply
39
model = LinearModuleForZeroStage3(16384, 16384)
40

41
model.to(get_accelerator().device_name()).half()
42

43
see_memory_usage("Before forward")
44
y = model(tens)
45

46
see_memory_usage("After forward")
47

48
model.weight.data = torch.zeros(1, dtype=torch.half, device=torch.device(get_accelerator().device_name()))
49

50
see_memory_usage("After weight zero")
51

52
y.backward(tens_back)
53