interleaver

interleaver

NNSIGHT_PREFIX module-attribute

NNSIGHT_PREFIX = '__nnsight'

Events

Bases: Enum

Enum for different types of events in the interleaving process.

VALUE class-attribute instance-attribute

VALUE = 'value'

SWAP class-attribute instance-attribute

SWAP = 'swap'

END class-attribute instance-attribute

END = 'end'

EXCEPTION class-attribute instance-attribute

EXCEPTION = 'exception'

SKIP class-attribute instance-attribute

SKIP = 'skip'

BARRIER class-attribute instance-attribute

BARRIER = 'barrier'

Cancelation

Bases: Exception

Exception raised when a request is canceled.

EarlyStopException

Bases: Exception

Exception raised to stop the execution of the model.

SkipException

SkipException(value: Any)

Bases: Exception

Exception raised to skip the execution of the model.

value instance-attribute

value = value

Interleaver

Interleaver(mediators: List[Mediator] = [], tracer: InterleavingTracer = None, batcher: Batcher = None)

Manages the interleaving of model execution and interventions.

This class coordinates the flow between the model's forward pass and user-defined intervention functions, allowing for inspection and modification of intermediate values.

ATTRIBUTE	DESCRIPTION
mediators	A dictionary of mediator names to mediator objects. Each meidator is responsible for a single invoke, or intervention function. TYPE: Dict[str, Mediator]
tracer	The tracer object that created this interleaver. Occationaly useful to know the tracer type for this interleaving. TYPE: Optional[InterleavingTracer]
batcher	The batcher object that manages the slice of inputs associtated with each mediator. TYPE: Batcher
current	The current mediator that is being processed. Must be update before resuming a given mediator. TYPE: Mediator

PARAMETER	DESCRIPTION
mediators	A list of mediator objects. TYPE: List[Mediator] DEFAULT: []
tracer	The tracer object that created this interleaver. TYPE: InterleavingTracer DEFAULT: None
batcher	The batcher object that manages the slice of inputs associtated with each mediator. TYPE: Batcher DEFAULT: None

interleaving property

interleaving: bool

Check if the interleaver is currently interleaving.

RETURNS	DESCRIPTION
bool	True if the interleaver is interleaving, False otherwise TYPE: bool

initialize

initialize(mediators: List[Mediator], tracer: InterleavingTracer, batcher: Batcher = None)

cancel

cancel()

Cancel all mediators / intervention threads.

iterate_provider

iterate_provider(provider: str)

Update a provider string to include which iteration of the provider is being provided.

PARAMETER	DESCRIPTION
provider	The provider string to update TYPE: str

RETURNS	DESCRIPTION
str	The updated provider string

Example

provider = "model.transformer.h[0].input" iterate_provider(provider) "model.transformer.h[0].input.i0"

provider = "model.transformer.h[0].input" iterate_provider(provider) "model.transformer.h[0].input.i1"

provider = "model.transformer.h[0].input" iterate_provider(provider) "model.transformer.h[0].input.i2"

iterate_requester

iterate_requester(requester: str)

Update a requester string to include which iteration of the requester is being requested. This is determined by the current mediator's iteration attribute, or influced by .iter contexts.

PARAMETER	DESCRIPTION
requester	The requester string to update TYPE: str

RETURNS	DESCRIPTION
str	The updated requester string

wrap_module

wrap_module(module: Module)

Instruments a PyTorch module to intercept inputs and outputs for interleaving.

PARAMETER	DESCRIPTION
module	The module to instrument TYPE: Module

RETURNS	DESCRIPTION
	None

wrap_operation

wrap_operation(fn: Callable, name: str, bound_obj: Optional[Any] = None)

Wrap an operation to intercept inputs and outputs for intervention, as well as the function itself. Used by Envoy.source to hook into intermediate operations of a forward pass.

PARAMETER	DESCRIPTION
fn	The intermediate operation function to wrap TYPE: Callable
name	The fully qualified name of the operation TYPE: str
bound_obj	The object fn is bound to if it is a method TYPE: Optional[Any] DEFAULT: None

RETURNS	DESCRIPTION
Callable	A wrapped version of the function

__enter__

__enter__()

__exit__

__exit__(exc_type, exc_val, exc_tb)

check_dangling_mediators

check_dangling_mediators()

check_cache_full

check_cache_full()

Print a warning if a module to be cached was missed.

handle

handle(provider: Optional[Any] = None, value: Optional[Any] = None, iterate: bool = False)

Handle a provider's value, allowing mediators to consume and modify it.

PARAMETER	DESCRIPTION
provider	The identifier of the provider TYPE: Optional[Any] DEFAULT: None
value	The value being provided TYPE: Optional[Any] DEFAULT: None
iterate	Whether to iterate the provider string TYPE: bool DEFAULT: False

RETURNS	DESCRIPTION
Any	The original or modified value

__deepcopy__

__deepcopy__(memo)

Mediator

Mediator(intervention: Callable, info: 'Tracer.Info', name: Optional[str] = None, batch_group: Optional[List[int]] = None, stop: Optional[int] = None)

Mediates between the model execution and a single intervention function.

This class handles the communication between the model's forward pass and a user-defined intervention function, allowing for inspection and modification of intermediate values.

ATTRIBUTE	DESCRIPTION
interleaver	The interleaver that this mediator is currently running in TYPE: Interleaver
intervention	The intervention function to mediate TYPE: Callable
info	Information about the tracing context associated with this mediator TYPE: Info
name	Optional name for the mediator TYPE: Optional[str]
batch_group	Optional batch group for the mediator to determine which slice of tensors are being intervened on TYPE: Optional[List[int]]
event_queue	Where the mediator (worker thread) puts events to be processed by the interleaver (main thread). Will only ever have 1 or 0 items in the queue. TYPE: SimpleQueue
response_queue	Where the interleaver (main thread) puts responses to events, to then be processed by the mediator (worker thread). Will only ever have 1 or 0 items in the queue. TYPE: SimpleQueue
worker	The thread that runs the intervention function TYPE: Thread
history	A set of providers that have been seen by the mediator. Used to detect out of order interventions. TYPE: Set[str]
iteration_tracker	A dictionary tracking the number of times each provider has been seen by the mediator. TYPE: Dict[str, int]
iteration	The current iteration this mediator is interventing on TYPE: int
user_cache	A list of caches to be used by the mediator TYPE: List[Cache]
all_stop	Optional number of times to execute this mediator TYPE: Optional[int]

PARAMETER	DESCRIPTION
intervention	The intervention function TYPE: Callable
info	Information about the tracing context TYPE: 'Tracer.Info'
name	Optional name for the mediator TYPE: Optional[str] DEFAULT: None
stop	Optional number of times to execute this mediator TYPE: Optional[int] DEFAULT: None

intervention instance-attribute

intervention = intervention

name instance-attribute

name = name if name else f'Mediator{id(self)}'

info instance-attribute

info = info

batch_group instance-attribute

batch_group = batch_group

interleaver instance-attribute

interleaver = None

event_queue instance-attribute

event_queue = Value()

response_queue instance-attribute

response_queue = Value()

worker instance-attribute

worker = None

history instance-attribute

history = set()

user_cache instance-attribute

user_cache: List['Cache'] = list()

iteration_tracker instance-attribute

iteration_tracker = defaultdict(int)

iteration instance-attribute

iteration = 0

all_stop instance-attribute

all_stop: Optional[int] = stop

args instance-attribute

args = list()

cross_invoker instance-attribute

cross_invoker = None

original_globals instance-attribute

original_globals = {}

alive property

alive

frame property

frame: FrameType

Get the frame of the intervention function.

RETURNS	DESCRIPTION
FrameType	The frame of the intervention function

OutOfOrderError

Bases: Exception

Exception raised when interventions are defined out of order.

Value

Value()

value instance-attribute

value = None

lock instance-attribute

lock = allocate_lock()

has_value instance-attribute

has_value = False

get

get()

wait

wait()

put

put(value: Any)

restore

restore(value: Any)

__enter__

__enter__()

__exit__

__exit__(exc_type, exc_val, exc_tb)

start

start(interleaver: Interleaver)

Start the mediator's intervention thread.

PARAMETER	DESCRIPTION
interleaver	The interleaver managing this mediator TYPE: Interleaver

cancel

cancel()

Cancel the intervention thread and its ephemeral state.

handle

handle(provider: Optional[str] = None)

Process a provider and its value. Depending on which event this mediator is waiting on, it will either: - Respond with the value - Swap (replace) the value with a new value - Respond with an out of order error - Skip the value - Set a barrier - End the execution (cancels the mediator)

PARAMETER	DESCRIPTION
provider	The identifier of the provider TYPE: Optional[str] DEFAULT: None

handle_value_event

handle_value_event(requester: Any, provider: Any) -> bool

Handle a value event by providing the requested value or recording a missed provider.

PARAMETER	DESCRIPTION
requester	The identifier of the requester TYPE: str
provider	The identifier of the provider TYPE: str

Returns: bool: Indicating whether the request was fulfilled by this processor, If so, continue processing events.

handle_swap_event

handle_swap_event(provider: Any, requester: Any, swap_value: Any)

Handle a swap event by swapping the value if the provider matches the requester.

PARAMETER	DESCRIPTION
requester	The identifier of the requester TYPE: str
provider	The identifier of the provider TYPE: str
swap_value	The value to swap in TYPE: Any

RETURNS	DESCRIPTION
bool	Indicating whether the swap was fulfilled by this processor, If so, continue processing events.

handle_exception_event

handle_exception_event(exception: Exception)

Handle an exception event by raising the exception.

PARAMETER	DESCRIPTION
exception	The exception to raise TYPE: Exception

RETURNS	DESCRIPTION
bool	Flag to stop processing events.

handle_barrier_event

handle_barrier_event(provider: Any, participants: Set[str])

Handle a barrier event by setting a barrier.

handle_end_event

handle_end_event()

Handle an end event by stopping the mediator.

handle_skip_event

handle_skip_event(provider: Any, requester: Any, value: Any)

respond

respond(value: Optional[Any] = None)

Respond from the interleaver (main thread) to the mediator (worker thread) the value for a pending event.

PARAMETER	DESCRIPTION
value	The value to provide TYPE: Optional[Any] DEFAULT: None

send

send(event: Events, requester: Any)

Send an event to interleaver (main thread) from this mediator (worker thread), and wait for it to be processed by the interleaver.

PARAMETER	DESCRIPTION
event	The event to send TYPE: Events
requester	The identifier of the requester, plus any additional data for the event. TYPE: Any

RETURNS	DESCRIPTION
Any	The response from the provider

request

request(requester: str)

Request a value from a specific provider.

PARAMETER	DESCRIPTION
requester	The identifier of the provider to request a value from TYPE: str

RETURNS	DESCRIPTION
Any	The requested value

swap

swap(requester: str, value: Any)

Send a swap event to replace the value of a provider.

PARAMETER	DESCRIPTION
requester	The identifier of the requester TYPE: str
value	The value to swap in TYPE: Any

stop

stop()

Stop the execution of the model by raising an EarlyStopException.

skip

skip(requester: Any, value: Any)

end

end()

Signal that execution should continue without further intervention.

exception

exception(exception: Exception)

Signal that an exception occurred during intervention.

PARAMETER	DESCRIPTION
exception	The exception that occurred TYPE: Exception

push

push()

Push local variables to the interleaver state.

pull

pull()

Pull variables from the interleaver state to the frame globals.

set_user_cache

set_user_cache(cache: 'Cache')

Set the user cache for this mediator.

PARAMETER	DESCRIPTION
cache	The cache to set TYPE: 'Cache'

__getstate__

__getstate__()

Get the state of the mediator for serialization.

__setstate__

__setstate__(state)

Set the state of the mediator for deserialization.