-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy path__init__.py
More file actions
77 lines (66 loc) · 2.78 KB
/
Copy path__init__.py
File metadata and controls
77 lines (66 loc) · 2.78 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
from supervisor import Context, Host
from typing import Dict, Sequence, TypedDict, Union, Optional, List
from typing_extensions import Unpack
import torch
from . import messages
from .history import RemoteException
from .stream import active_stream, Stream
from .future import Future
from .tensor import Tensor, dtensor_check
from .device_mesh import DeviceMesh, remote_function, active_mesh, RemoteFunction
from .ndslice import NDSlice
from . import stream, device_mesh
from .controller import Controller as _Controller
from .backend import ProcessBackend as _ProcessBackend
def fetch_shard(obj, coordinates: Optional[Dict[str, int]] = None, preprocess: Optional[RemoteFunction] = None):
"""
Retrieve the shard at `coordinates` of the current device mesh of each tensor in obj.
obj - a pytree containing the tensors the fetch
coordinates - a dictionary from mesh dimension name to coordinate of the shard
If None, this will fetch from coordinate 0 for all dimensions (useful after all_reduce/all_gather)
preprocess - an optional specifier for a remote function that is applyied to obj before returning the value.
This can be used to turn the tensor into some non-tensor values before copying it back
"""
if device_mesh._active is None:
raise RuntimeError("No device mesh active")
mesh = device_mesh._active
ctrl = mesh.ctrl
if preprocess is None:
remote = RemoteFunction('fetch_shard', lambda x: x)
else:
remote = preprocess
fake_result, dtensors, mutates = dtensor_check(remote, (obj,), {}, mesh, stream._active)
fut = Future(ctrl)
ident = ctrl.history.ident(mutates, dtensors, fut)
process = mesh._process(coordinates)
ctrl.send(process, messages.FetchValue(ident, mutates, preprocess, obj, stream._active))
return fut
class TensorOptions(TypedDict, total=False):
dtype: 'torch.dtype'
layer: 'torch.layout'
device: Union[str, 'torch.device']
requires_grad: bool
pin_memory: bool
memory_format: 'torch.memory_format'
class Pipe:
def push(self, tensor: 'Tensor'):
raise NotImplementedError()
def pop(self, sizes: Sequence[int], **kwargs: Unpack[TensorOptions]) -> 'Tensor':
raise NotImplementedError()
def world_mesh(ctx: Context, hosts: List[Host], gpu_per_host: int, _processes=None):
backend = _ProcessBackend(ctx, hosts, gpu_per_host, _processes=_processes)
ctrl = _Controller(backend)
return DeviceMesh(ctrl, NDSlice(0, [len(hosts), gpu_per_host], [gpu_per_host, 1]), ('host', 'gpu'))
def get_active_stream():
return stream._active
__all__ = [
'RemoteException',
'remote_function',
'Future',
'active_stream',
'active_mesh',
'DeviceMesh',
'Tensor',
'Stream',
'get_active_stream',
]