rlightning.weights.weight_buffer

class rlightning.weights.weight_buffer.CPUWeightBuffer(config: WeightBufferConfig)

Bases: WeightBuffer

Implements a CPU buffer strategy. Weights are transferred to and stored in CPU memory on the evaluator nodes and moved to the GPU just-in-time for model updates.

VALUE_TYPE: Type = MockModule('numpy.ndarray')

class rlightning.weights.weight_buffer.ShardedWeightBuffer(config: WeightBufferConfig)

Bases: WeightBuffer

Implements a sharded buffer strategy for memory efficiency. Each GPU holds only a fraction of the total weights. Weights are gathered across GPUs on the same node just-in-time for model updates.

add_shard(dtype: MockModule('torch.dtype'), flat_tensor_shard: MockModule('torch.Tensor'), metadata_list: List[Dict], numel_total: int, is_last: bool)

Adds a weight shard for a specific dtype to the buffer.

apply_to_model(policy)

Called by ShardedWeightBuffer. This reconstructs the full weights from shards and applies them to the model.

clear()

class rlightning.weights.weight_buffer.WeightBuffer(config: WeightBufferConfig)

Bases: object

policy weight buffer, for update weights from train policy to eval policy

VALUE_TYPE: Type = MockModule('torch.Tensor')

add(module_name: str, name: str, weight: MockModule('torch.Tensor'), is_last: bool = False)

add_state_dict(state_dict)

Add a complete state dictionary to the buffer.

clear()

get_state_dict()

Retrieves the state dictionary from the buffer.

is_ready()

preprocess(*args, **kwargs)

sample()