218 lines
8.4 KiB
Python
218 lines
8.4 KiB
Python
# Copyright (c) Kyutai, all rights reserved.
|
|
# This source code is licensed under the license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
"""
|
|
Streaming module API that should be implemented by all Streaming components,
|
|
"""
|
|
|
|
import abc
|
|
from contextlib import ExitStack
|
|
from dataclasses import dataclass
|
|
import typing as tp
|
|
from torch import nn
|
|
import torch
|
|
from ..utils.compile import CUDAGraphed
|
|
|
|
|
|
@dataclass
|
|
class State(abc.ABC):
|
|
"""Base State for streaming, requires to be resetable and also support the context
|
|
protocol. The state __enter__ and __exit__ will be called upon entering and exiting
|
|
streaming in the parent module, but not upon `reset()` calls.
|
|
"""
|
|
|
|
batch_size: int
|
|
device: torch.device
|
|
|
|
def __post_init__(self):
|
|
self.exec_mask = torch.ones(self.batch_size, dtype=torch.bool, device=self.device)
|
|
self._set_exec_mask_graphed: CUDAGraphed | None = None
|
|
|
|
def set_exec_mask(self, exec_mask: torch.Tensor):
|
|
self.exec_mask[:] = exec_mask
|
|
|
|
def reset(self, reset_mask: torch.Tensor) -> None:
|
|
self.exec_mask[:] = torch.where(reset_mask, torch.ones_like(self.exec_mask), self.exec_mask)
|
|
|
|
def __enter__(self) -> None:
|
|
pass
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback) -> None:
|
|
pass
|
|
|
|
|
|
StateT = tp.TypeVar("StateT", bound=State)
|
|
|
|
|
|
class StreamingModule(abc.ABC, nn.Module, tp.Generic[StateT]):
|
|
"""Common API for streaming components.
|
|
|
|
Each streaming component has a streaming state, `self._streaming_state`, which is None by default.
|
|
|
|
To set a streaming component in streaming state, use
|
|
|
|
with module.streaming(batch_size):
|
|
...
|
|
|
|
This will automatically void the streaming state when exiting the context manager.
|
|
This also automatically propagates to all streaming children module.
|
|
When the streaming state is set, modules should store whatever state they need in there.
|
|
"""
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self._streaming_state: StateT | None = None
|
|
self._streaming_detached: bool = False
|
|
self._cached_children: list[tuple[str, StreamingModule]] | None = None
|
|
|
|
@property
|
|
def is_streaming(self):
|
|
return self._streaming_state is not None
|
|
|
|
def set_streaming_detached(self, streaming_detached: bool):
|
|
"""If set to False, the default, this module and all submodules will switch to streaming mode
|
|
if a parent module is set to streaming mode.
|
|
If set to True, or in detach mode, only a direct call to this module `.streaming(...)` method
|
|
will set it into streaming mode, ignoring the changes from its parents.
|
|
|
|
This is useful if streaming over two different dimensions, e.g. for the RQ-Transformer
|
|
with the inner Depth Transformer working on the dimension of the codebooks."""
|
|
self._streaming_detached = streaming_detached
|
|
|
|
def _apply_named_streaming(self, fn: tp.Any):
|
|
def _handle_module(prefix: str, module: nn.Module):
|
|
if isinstance(module, StreamingModule):
|
|
# If prefix is empty, we are the direct receiver of the streaming request,
|
|
# otherwise, we are inheriting from a parent and will stop if detached.
|
|
if module._streaming_detached and prefix != "":
|
|
return
|
|
assert self._cached_children is not None
|
|
self._cached_children.append((prefix, module))
|
|
for name, child in module.named_children():
|
|
if prefix:
|
|
new_prefix = prefix + "." + name
|
|
else:
|
|
new_prefix = name
|
|
_handle_module(new_prefix, child)
|
|
|
|
if self._cached_children is None:
|
|
self._cached_children = []
|
|
_handle_module("", self)
|
|
for name, child in self._cached_children:
|
|
fn(name, child)
|
|
|
|
def _start_streaming(self, batch_size: int, exit_stack: ExitStack):
|
|
def _start_streaming(name: str, module: StreamingModule):
|
|
assert module._streaming_state is None, f"{name} is already streaming!"
|
|
state = module._init_streaming_state(batch_size)
|
|
exit_stack.enter_context(state)
|
|
module._streaming_state = state
|
|
|
|
self._apply_named_streaming(_start_streaming)
|
|
|
|
def _stop_streaming(self) -> None:
|
|
def _stop_streaming(name: str, module: StreamingModule):
|
|
module._streaming_state = None
|
|
|
|
self._apply_named_streaming(_stop_streaming)
|
|
|
|
@abc.abstractmethod
|
|
def _init_streaming_state(self, batch_size: int) -> StateT: ...
|
|
|
|
def streaming_forever(self, batch_size: int):
|
|
self.streaming(batch_size).__enter__()
|
|
|
|
def streaming(self, batch_size: int) -> ExitStack:
|
|
"""Context manager to enter streaming mode. Reset streaming state on exit."""
|
|
|
|
exit_stack = ExitStack()
|
|
self._start_streaming(batch_size, exit_stack)
|
|
exit_stack.callback(self._stop_streaming)
|
|
return exit_stack
|
|
|
|
def reset_streaming(self, reset_mask: torch.Tensor | None = None) -> None:
|
|
"""Reset the streaming state."""
|
|
|
|
def _reset(name: str, module: StreamingModule):
|
|
state = module._streaming_state
|
|
if state is None:
|
|
raise ValueError(
|
|
f"Trying to reset streaming, but {name} wasn't streaming."
|
|
)
|
|
state.reset(reset_mask)
|
|
|
|
state = self._streaming_state
|
|
assert state is not None
|
|
if reset_mask is None:
|
|
reset_mask = torch.ones(state.batch_size, device=state.device, dtype=torch.bool)
|
|
else:
|
|
reset_mask = reset_mask.to(state.device)
|
|
self._apply_named_streaming(_reset)
|
|
|
|
def get_streaming_state(self) -> dict[str, tp.Any]:
|
|
"""Return the complete streaming state, including that of sub-modules."""
|
|
state: dict[str, tp.Any] = {}
|
|
|
|
def _add(name: str, module: StreamingModule):
|
|
state[name] = module._streaming_state
|
|
|
|
self._apply_named_streaming(_add)
|
|
return state
|
|
|
|
def set_streaming_state(self, state: dict[str, tp.Any]):
|
|
"""Set the streaming state, including that of sub-modules."""
|
|
state = dict(state)
|
|
|
|
def _set(name: str, module: StreamingModule):
|
|
if name in state:
|
|
module._streaming_state = state[name]
|
|
state.pop(name)
|
|
else:
|
|
raise RuntimeError(f"Expected to find a streaming state for {name}.")
|
|
|
|
self._apply_named_streaming(_set)
|
|
if state:
|
|
raise RuntimeError(f"Some states were not consumed: {list(state.keys())}")
|
|
|
|
def set_exec_mask(self, exec_mask: torch.Tensor):
|
|
"""Set the execution mask, a tensor of boolean of shape `(B,), indicating
|
|
for each batch item whether the internal state should be updated or not as if
|
|
real data had been received.
|
|
|
|
This is useful for running desynchronized streams with batching, e.g. when
|
|
the mask is False for an entry, the internal state will be unchanged by the provided
|
|
data, e.g. will be on the next step as if the previous one had never happened.
|
|
|
|
There is no magic here, each StreamingModule subclass is responsible for respecting
|
|
the exec_mask.
|
|
"""
|
|
|
|
state = self._streaming_state
|
|
assert state is not None
|
|
exec_mask = exec_mask.to(state.device)
|
|
|
|
def _set_exec_mask(exec_mask: torch.Tensor):
|
|
def _set_exec_mask_fn(name: str, module: StreamingModule):
|
|
state = module._streaming_state
|
|
assert state is not None
|
|
state.set_exec_mask(exec_mask)
|
|
self._apply_named_streaming(_set_exec_mask_fn)
|
|
|
|
if state._set_exec_mask_graphed is None:
|
|
disable = state.device.type != 'cuda'
|
|
state._set_exec_mask_graphed = CUDAGraphed(_set_exec_mask, disable=disable)
|
|
|
|
state._set_exec_mask_graphed(exec_mask)
|
|
|
|
|
|
class StreamingContainer(StreamingModule[State]):
|
|
def _init_streaming_state(self, batch_size: int) -> State:
|
|
device = next(iter(self.parameters())).device
|
|
return State(batch_size, device)
|