"...python/git@developer.sourcefind.cn:zhaoyu6/sglang.git" did not exist on "d738ab52f86f65806fc6549270a704e34aae5a32"
test.py 1.59 KB
Newer Older
aiss's avatar
aiss committed
1
2
'''Copyright The Microsoft DeepSpeed Team'''

Samyam Rajbhandari's avatar
Samyam Rajbhandari committed
3
4
5
import torch
from deepspeed.pt.deepspeed_linear import LinearModuleForZeroStage3
from deepspeed.pt.log_utils import logger
aiss's avatar
aiss committed
6
from deepspeed.accelerator import get_accelerator
Samyam Rajbhandari's avatar
Samyam Rajbhandari committed
7
8
9
10
11
12
13
14


def see_memory_usage(message):

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


aiss's avatar
aiss committed
31
32
33
34
tens = torch.rand(1024,
                  16384,
                  dtype=torch.half,
                  device=torch.device(get_accelerator().device_name()))
Samyam Rajbhandari's avatar
Samyam Rajbhandari committed
35
36
37
38
39
40
tens_back = tens.detach().clone()

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

aiss's avatar
aiss committed
41
model.to(get_accelerator().device_name()).half()
Samyam Rajbhandari's avatar
Samyam Rajbhandari committed
42
43
44
45
46
47

see_memory_usage("Before forward")
y = model(tens)

see_memory_usage("After forward")

aiss's avatar
aiss committed
48
49
50
model.weight.data = torch.zeros(1,
                                dtype=torch.half,
                                device=torch.device(get_accelerator().device_name()))
Samyam Rajbhandari's avatar
Samyam Rajbhandari committed
51
52
53
54

see_memory_usage("After weight zero")

y.backward(tens_back)