wehub-resource-sync
77 lines
3.0 KiB
Python
77 lines
3.0 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import torch
|
|
|
|
try:
|
|
from torch._subclasses import FakeTensorMode
|
|
from torch._subclasses.fake_tensor import unset_fake_temporarily
|
|
from torch._dynamo.variables.builder import wrap_to_fake_tensor_and_record
|
|
except ImportError:
|
|
# Unsupported torch version
|
|
pass
|
|
|
|
|
|
def wrap_if_ds_param(t):
|
|
if hasattr(t, 'ds_id'):
|
|
data = torch.rand(t.ds_shape,
|
|
dtype=t.dtype,
|
|
layout=t.layout,
|
|
device=t.device,
|
|
pin_memory=t.is_pinned(),
|
|
requires_grad=t.requires_grad)
|
|
if isinstance(t, torch.nn.Parameter):
|
|
t = torch.nn.Parameter(data, requires_grad=t.requires_grad)
|
|
else:
|
|
t = data
|
|
return t
|
|
|
|
|
|
def _get_guard_sizes_strides(t):
|
|
if hasattr(t, "ds_id"):
|
|
# ZeRO-3 may temporarily all-gather a parameter during tracing, but the
|
|
# stable module state used by TorchDynamo guards is the released
|
|
# partitioned form, where DeepSpeed resets param.data to empty(0).
|
|
released = torch.empty(0, dtype=t.dtype, device=t.device)
|
|
return released.size(), released.stride()
|
|
|
|
return t.size(), t.stride()
|
|
|
|
|
|
def patch_fake_tensor():
|
|
# dynamo tracer uses wrap_to_fake_tensor_and_record
|
|
# Wrapping FakeTensorMode.from_tensor is not sufficient as dynamo generates SymbolicContext before calling from_tensor
|
|
original_wrap_to_fake_tensor_and_record = wrap_to_fake_tensor_and_record
|
|
|
|
def wrap_to_fake_tensor_and_record_wrapper(t, *args, **kwargs):
|
|
dummy_tensor = wrap_if_ds_param(t)
|
|
ret = original_wrap_to_fake_tensor_and_record(dummy_tensor, *args, **kwargs)
|
|
tx = kwargs.get("tx") if "tx" in kwargs else args[0]
|
|
source = kwargs.get("source")
|
|
if tracing_context := torch._guards.TracingContext.try_get():
|
|
tracing_context.tensor_to_context[t] = tracing_context.tensor_to_context.pop(dummy_tensor)
|
|
if source is not None:
|
|
# Keep the full ds_shape symbolic context from the dummy tensor, but
|
|
# use the stable released ZeRO-3 parameter representation for
|
|
# TorchDynamo's tensor-match guards. PyTorch 2.9 started enforcing
|
|
# those guards for parameters during build_guards().
|
|
size, stride = _get_guard_sizes_strides(t)
|
|
tx.output.input_source_to_sizes_strides[source] = {
|
|
"size": size,
|
|
"stride": stride,
|
|
}
|
|
return ret
|
|
|
|
torch._dynamo.variables.builder.wrap_to_fake_tensor_and_record = wrap_to_fake_tensor_and_record_wrapper
|
|
|
|
# aot_module_simplified uses fake_mode.from_tensor to process inputs
|
|
original_from_tensor = FakeTensorMode.from_tensor
|
|
|
|
def from_tensor_wrapper(self, t, *args, **kwargs):
|
|
with unset_fake_temporarily():
|
|
return original_from_tensor(self, wrap_if_ds_param(t), *args, **kwargs)
|
|
|
|
FakeTensorMode.from_tensor = from_tensor_wrapper
|