SYMBOL INDEX (2598 symbols across 253 files) FILE: examples/custom_pipeline_example/custom_grpo.py function example_builtin_grpo (line 33) | def example_builtin_grpo() -> TaskGraph: function grpo_with_custom_reward (line 48) | def grpo_with_custom_reward() -> TaskGraph: function my_custom_reward_fn (line 97) | def my_custom_reward_fn(batch: TensorDict, **kwargs) -> NodeOutput: FILE: examples/custom_reward/rewardfunc_gsm8k.py function extract_solution (line 4) | def extract_solution(solution_str, method="strict"): function compute_score (line 30) | def compute_score(data_sources, solution_strs, ground_truths, extra_info... FILE: examples/data_preprocess/deepscaler.py function load_json (line 28) | def load_json(file_path): function make_map_fn (line 55) | def make_map_fn(split_name): FILE: examples/data_preprocess/geo3k.py function make_map_fn (line 45) | def make_map_fn(split): FILE: examples/data_preprocess/gsm8k.py function extract_solution (line 27) | def extract_solution(solution_str): function make_map_fn (line 52) | def make_map_fn(split): FILE: examples/data_preprocess/math_dataset.py function extract_solution (line 27) | def extract_solution(solution_str): function make_map_fn (line 50) | def make_map_fn(split): FILE: examples/data_preprocess/mm_eureka.py function make_map_fn (line 53) | def make_map_fn(split): FILE: examples/experimental/marft/config/code_env.py class CodeEnv (line 18) | class CodeEnv(): method __init__ (line 19) | def __init__(self): method reset (line 21) | def reset(self) -> Any: method step (line 23) | async def step(self, actions, ground_truth): FILE: examples/experimental/marft/config/math_env.py class MathEnv (line 18) | class MathEnv(): method __init__ (line 19) | def __init__(self): method reset (line 21) | def reset(self) -> Any: method step (line 23) | async def step(self, actions, ground_truth): FILE: examples/experimental/marft/config/process.py function pre_process (line 16) | def pre_process(tokenizer, prompt_id, obs, **kwargs): function post_process (line 26) | def post_process(tokenizer, prompt_id, response_id, **kwargs): FILE: siirl/dag_worker/checkpoint_manager.py class CheckpointManager (line 30) | class CheckpointManager: method __init__ (line 33) | def __init__( method save_checkpoint (line 53) | def save_checkpoint(self, global_steps: int) -> None: method _save_model_states (line 73) | def _save_model_states(self, global_steps: int, step_dir: str) -> None: method _save_dataloader_state (line 111) | def _save_dataloader_state(self, step_dir: str) -> None: method _commit_checkpoint (line 124) | def _commit_checkpoint(self, global_steps: int) -> None: method load_checkpoint (line 134) | def load_checkpoint(self) -> int: method _determine_checkpoint_path (line 172) | def _determine_checkpoint_path(self) -> Optional[str]: method _load_model_states (line 199) | def _load_model_states(self, global_step_folder: str) -> None: method _load_dataloader_state (line 237) | def _load_dataloader_state(self, global_step_folder: str) -> None: FILE: siirl/dag_worker/constants.py class DAGInitializationError (line 19) | class DAGInitializationError(Exception): class DAGConstants (line 25) | class DAGConstants: FILE: siirl/dag_worker/core_algos.py function register_policy_loss (line 54) | def register_policy_loss(name: str) -> Callable[[PolicyLossFn], PolicyLo... function get_policy_loss_fn (line 71) | def get_policy_loss_fn(name): function compute_response_mask (line 89) | def compute_response_mask(data: TensorDict): function register_adv_est (line 130) | def register_adv_est(name_or_enum: str | AdvantageEstimator) -> Any: function get_adv_estimator_fn (line 151) | def get_adv_estimator_fn(name_or_enum): class AdaptiveKLController (line 167) | class AdaptiveKLController: method __init__ (line 173) | def __init__(self, init_kl_coef, target_kl, horizon): method update (line 178) | def update(self, current_kl, n_steps): class FixedKLController (line 191) | class FixedKLController: method __init__ (line 194) | def __init__(self, kl_coef): method update (line 197) | def update(self, current_kl, n_steps): function get_kl_controller (line 207) | def get_kl_controller(kl_ctrl): function compute_gae_advantage_return (line 230) | def compute_gae_advantage_return( function compute_grpo_outcome_advantage (line 280) | def compute_grpo_outcome_advantage( function compute_marft_gae_advantage_return (line 356) | def compute_marft_gae_advantage_return( function compute_cpgd_outcome_advantage (line 466) | def compute_cpgd_outcome_advantage( function compute_rewards (line 525) | def compute_rewards(token_level_scores, old_log_prob, ref_log_prob, kl_r... function agg_loss (line 541) | def agg_loss(loss_mat: torch.Tensor, loss_mask: torch.Tensor, loss_agg_m... function compute_policy_loss_cpgd (line 578) | def compute_policy_loss_cpgd( function compute_policy_loss (line 636) | def compute_policy_loss( function compute_policy_loss_vanilla (line 716) | def compute_policy_loss_vanilla( function compute_policy_loss_gspo (line 801) | def compute_policy_loss_gspo( function compute_policy_loss_gpg (line 871) | def compute_policy_loss_gpg( function compute_policy_loss_clip_cov (line 904) | def compute_policy_loss_clip_cov( function compute_policy_loss_kl_cov (line 1004) | def compute_policy_loss_kl_cov( function compute_policy_loss_geo_mean (line 1080) | def compute_policy_loss_geo_mean( function compute_entropy_loss (line 1161) | def compute_entropy_loss(logits, response_mask, loss_agg_mode: str = "to... function compute_value_loss (line 1178) | def compute_value_loss( function kl_penalty (line 1220) | def kl_penalty(logprob: torch.FloatTensor, ref_logprob: torch.FloatTenso... function kl_penalty_forward (line 1246) | def kl_penalty_forward(logprob: torch.FloatTensor, ref_logprob: torch.Fl... function compute_pf_ppo_reweight_data (line 1284) | def compute_pf_ppo_reweight_data( function apply_kl_penalty (line 1363) | def apply_kl_penalty(data: TensorDict, kl_ctrl: AdaptiveKLController, kl... function compute_advantage (line 1405) | def compute_advantage(data: TensorDict, adv_estimator, gamma=1.0, lam=1.... FILE: siirl/dag_worker/dag_utils.py function timer (line 57) | def timer(enable_perf: bool, name: str, timing_dict: dict): function add_prefix_to_dataproto (line 69) | def add_prefix_to_dataproto(tensordict: TensorDict, node: Node): function remove_prefix_from_dataproto (line 95) | def remove_prefix_from_dataproto(tensordict, node: Node): function add_prefix_to_metrics (line 121) | def add_prefix_to_metrics(metrics: dict, node: Node) -> dict: function get_and_validate_rank (line 141) | def get_and_validate_rank() -> int: function get_taskgraph_for_rank (line 152) | def get_taskgraph_for_rank(rank: int, taskgraph_mapping: Dict[int, TaskG... function log_ray_actor_info (line 164) | def log_ray_actor_info(rank: int): function log_role_worker_mapping (line 176) | def log_role_worker_mapping(role_worker_mapping: Dict[NodeRole, Type[Wor... function find_first_non_compute_ancestor (line 195) | def find_first_non_compute_ancestor(taskgraph: TaskGraph, start_node_id:... function should_create_worker (line 228) | def should_create_worker(role_worker_mapping: Dict[NodeRole, Type[Worker... function generate_node_worker_key (line 236) | def generate_node_worker_key(node: Node) -> str: function setup_sharding_manager (line 241) | def setup_sharding_manager( function get_worker_classes (line 370) | def get_worker_classes(config, strategy: str) -> Dict[NodeRole, Type[Wor... function get_parallelism_config (line 414) | def get_parallelism_config(reference_node: Node) -> tuple[int, int]: function prepare_generation_batch (line 450) | def prepare_generation_batch(batch: TensorDict) -> TensorDict: function prepare_local_batch_metrics (line 465) | def prepare_local_batch_metrics(batch: TensorDict, use_critic: bool = Tr... function whether_put_data (line 515) | def whether_put_data(rank, is_current_last_pp_tp_rank0, next_dp_size, cu... function reduce_and_broadcast_metrics (line 542) | def reduce_and_broadcast_metrics( function format_metrics_by_group (line 580) | def format_metrics_by_group(metrics: Dict[str, Any], group_order: List[s... function log_metrics_to_console (line 629) | def log_metrics_to_console(rank: int, ordered_metrics: List[Tuple[str, A... function dump_validation_generations (line 638) | def dump_validation_generations( function aggregate_and_write_performance_metrics (line 684) | def aggregate_and_write_performance_metrics( function log_core_performance_metrics (line 775) | def log_core_performance_metrics(rank: int, enable_perf: bool, metrics: ... function get_time_now (line 884) | def get_time_now(time_zone: str = "Asia/Shanghai") -> datetime: function consistent_hash (line 889) | def consistent_hash(s: str) -> int: FILE: siirl/dag_worker/dagworker.py class DAGWorker (line 81) | class DAGWorker(Worker): method __init__ (line 87) | def __init__( method execute_task_graph (line 150) | def execute_task_graph(self): method _run_training_loop (line 174) | def _run_training_loop(self): method _cleanup_step_buffers (line 256) | def _cleanup_step_buffers(self, timing_raw: dict) -> None: method _run_training_step (line 272) | def _run_training_step(self, epoch: int, batch_idx: int) -> Optional[L... method generate_sync_mode (line 469) | def generate_sync_mode(self, agent_group, batch: TensorDict) -> NodeOu... method generate_async_mode (line 478) | def generate_async_mode(self, batch: TensorDict) -> NodeOutput: method generate_multi_agent_mode (line 490) | def generate_multi_agent_mode(self, config, batch: TensorDict) -> Node... method generate_embodied_mode (line 508) | def generate_embodied_mode(self, agent_group, batch: TensorDict, **kwa... method generate (line 557) | def generate(self, config, batch: TensorDict, **kwargs) -> NodeOutput: method compute_reward (line 575) | def compute_reward(self, config, batch: TensorDict, **kwargs) -> NodeO... method compute_old_log_prob (line 604) | def compute_old_log_prob(self, config, batch: TensorDict, **kwargs) ->... method compute_ref_log_prob (line 624) | def compute_ref_log_prob(self, config, batch: TensorDict, **kwargs) ->... method compute_value (line 632) | def compute_value(self, config, batch: TensorDict, **kwargs) -> NodeOu... method compute_multi_agent_advantage (line 639) | def compute_multi_agent_advantage(self, config, batch: TensorDict, **k... method compute_advantage (line 686) | def compute_advantage(self, config, batch: TensorDict, **kwargs) -> No... method train_critic (line 708) | def train_critic(self, config, batch: TensorDict, **kwargs) -> NodeOut... method train_actor (line 716) | def train_actor(self, config, batch: TensorDict, **kwargs) -> NodeOutput: method _initialize_worker (line 732) | def _initialize_worker(self): method _setup_distributed_environment (line 767) | def _setup_distributed_environment(self): method _setup_tokenizers (line 817) | def _setup_tokenizers(self): method _setup_dataloader (line 843) | def _setup_dataloader(self): method _setup_reward_managers (line 877) | def _setup_reward_managers(self): method _setup_role_worker_mapping (line 893) | def _setup_role_worker_mapping(self): method _initialize_node_workers (line 913) | def _initialize_node_workers(self): method init_graph (line 965) | def init_graph(self): method _load_model_weights (line 996) | def _load_model_weights(self): method _setup_sharding_manager (line 1018) | def _setup_sharding_manager(self): method _setup_async_rollout (line 1037) | def _setup_async_rollout(self): method _setup_multi_agent_loop (line 1052) | def _setup_multi_agent_loop(self): method _init_validator (line 1071) | def _init_validator(self): method _init_metrics_collector (line 1095) | def _init_metrics_collector(self): method _init_checkpoint_manager (line 1112) | def _init_checkpoint_manager(self): method init_async_server (line 1126) | def init_async_server(self, node:Node, node_worker): method put_data_to_buffers (line 1150) | def put_data_to_buffers( method get_data_from_buffers (line 1221) | def get_data_from_buffers( method reset_data_buffer (line 1307) | def reset_data_buffer(self): method _get_node_process_group (line 1317) | def _get_node_process_group(self, node: Node) -> ProcessGroup: method _get_node_dp_info (line 1328) | def _get_node_dp_info(self, node: Node) -> tuple[int, int, int, int, i... method get_zeromq_address (line 1376) | def get_zeromq_address(self): method multi_agent_put_log (line 1379) | def multi_agent_put_log(self, key: str, data: TensorDict, agent_group:... method check_mode (line 1384) | def check_mode(self): FILE: siirl/dag_worker/data_structures.py class ValidationResult (line 24) | class ValidationResult: class ValidationPayload (line 36) | class ValidationPayload: class NodeOutput (line 46) | class NodeOutput: FILE: siirl/dag_worker/metric_aggregator.py class _ReduceOp (line 25) | class _ReduceOp(Enum): class DistributedMetricAggregator (line 54) | class DistributedMetricAggregator: method __init__ (line 60) | def __init__( method _bucket_local_metrics (line 78) | def _bucket_local_metrics(self, metrics: Dict, expected_keys: set = No... method aggregate_and_get_results (line 173) | def aggregate_and_get_results(self) -> Dict[str, float]: FILE: siirl/dag_worker/metrics_collector.py class MetricsCollector (line 50) | class MetricsCollector: method __init__ (line 63) | def __init__( method collect_final_metrics (line 96) | def collect_final_metrics(self, batch: TensorDict, timing_raw: dict) -... method collect_multi_agent_final_metrics (line 240) | def collect_multi_agent_final_metrics(self, batch: TensorDict, ordered... FILE: siirl/dag_worker/validator.py class Validator (line 43) | class Validator: method __init__ (line 57) | def __init__( method validate (line 122) | def validate(self, global_step: int) -> Dict[str, float]: method _validate_embodied (line 131) | def _validate_embodied(self, global_step) -> Dict[str, float]: method _validate_text_generation (line 200) | def _validate_text_generation(self, global_step: int) -> Dict[str, flo... method _generate_for_validation (line 264) | def _generate_for_validation(self, batch: TensorDict) -> TensorDict: method _score_and_package_results (line 306) | def _score_and_package_results(self, generated_proto: TensorDict) -> L... method _aggregate_and_log_validation_metrics (line 359) | def _aggregate_and_log_validation_metrics(self, all_payloads: List[Val... method _aggregate_validation_results (line 395) | def _aggregate_validation_results(self, all_payloads: List[ValidationP... method _prepare_embodied_validation_batch (line 459) | def _prepare_embodied_validation_batch(self) -> TensorDict: method _generate_for_embodied_validation (line 475) | def _generate_for_embodied_validation(self, batch: TensorDict, global_... method _score_embodied_results (line 519) | def _score_embodied_results(self, generated_proto: TensorDict) -> List... method _aggregate_and_log_embodied_metrics (line 580) | def _aggregate_and_log_embodied_metrics(self, all_payloads: List[Valid... method _aggregate_embodied_results (line 622) | def _aggregate_embodied_results(self, all_payloads: List[ValidationPay... FILE: siirl/data_coordinator/data_buffer.py class DataCoordinator (line 29) | class DataCoordinator: method __init__ (line 36) | def __init__(self, nnodes: int, ppo_mini_batch_size: int, world_size: ... method put (line 56) | async def put(self, sample_info: SampleInfo, sample_ref: Any, caller_n... method put_batch (line 99) | async def put_batch(self, sample_infos: List[SampleInfo], sample_refs:... method get_batch (line 137) | async def get_batch( method _log_accumulation_progress (line 219) | def _log_accumulation_progress(self, current_samples: int, target_samp... method _log_dispatch_stats (line 242) | def _log_dispatch_stats(self, dispatched_samples: int): method _apply_length_balancing (line 268) | def _apply_length_balancing( method _apply_length_balancing_single_sample (line 353) | def _apply_length_balancing_single_sample( method get_all_by_filter (line 416) | async def get_all_by_filter(self, filter_plugin: Callable[[SampleInfo]... method get_valid_size (line 438) | async def get_valid_size(self) -> int: method peek_source_dp_size (line 443) | async def peek_source_dp_size(self, filter_plugin: Callable[[SampleInf... method reset_cache (line 461) | def reset_cache(self): method __repr__ (line 468) | def __repr__(self) -> str: function init_data_coordinator (line 476) | def init_data_coordinator(num_buffers: int, ppo_mini_batch_size: int, wo... FILE: siirl/data_coordinator/dataloader/data_loader_node.py class RepeatDataset (line 29) | class RepeatDataset(torch.utils.data.Dataset): method __init__ (line 54) | def __init__(self, base_dataset, repeat_factor): method __len__ (line 59) | def __len__(self): method __getitem__ (line 62) | def __getitem__(self, idx): class DataLoaderNode (line 66) | class DataLoaderNode(Node): method __init__ (line 73) | def __init__( method _create_dataloader (line 135) | def _create_dataloader(self): method _reinit_dataloader_sampler (line 246) | def _reinit_dataloader_sampler(self): method get_train_dataloader (line 274) | def get_train_dataloader(self): method get_val_dataloader (line 283) | def get_val_dataloader(self): method run (line 292) | def run(self, epoch: Optional[int] = None, is_validation_step: bool = ... method state_dict (line 407) | def state_dict(self) -> Dict[str, Any]: method load_state_dict (line 418) | def load_state_dict(self, state_dict: Dict[str, Any]): FILE: siirl/data_coordinator/dataloader/embodied_preprocess.py function prepare_libero_train_valid_datasets (line 23) | def prepare_libero_train_valid_datasets( FILE: siirl/data_coordinator/dataloader/partitioned_dataset.py function collate_fn (line 37) | def collate_fn(data_list: list[dict]) -> dict: class PartitionedRLHFDataset (line 73) | class PartitionedRLHFDataset(Dataset): method __init__ (line 97) | def __init__( method _load_partitioned_raw_data (line 211) | def _load_partitioned_raw_data(self, dataset_files: Sequence[str]) -> ... method _filter_overlong_prompts (line 297) | def _filter_overlong_prompts(self, raw_dataframe: datasets.Dataset) ->... method __len__ (line 310) | def __len__(self) -> int: method _build_messages (line 313) | def _build_messages(self, example: dict) -> list: method _preprocess_function (line 333) | def _preprocess_function(self, row_dict: Dict) -> Dict: method __getitem__ (line 444) | def __getitem__(self, item: int) -> Dict: FILE: siirl/data_coordinator/dataloader/vision_utils.py function process_image (line 23) | def process_image(image: Union[dict, Image.Image], max_pixels: int, min_... function process_video (line 74) | def process_video( function process_multi_modal_inputs_for_minicpmo (line 113) | def process_multi_modal_inputs_for_minicpmo(input_ids, attention_mask, p... FILE: siirl/data_coordinator/protocol.py class _TensorDictConfigMeta (line 49) | class _TensorDictConfigMeta(type): method auto_padding (line 55) | def auto_padding(cls): method auto_padding (line 60) | def auto_padding(cls, enabled: bool): class TensorDictConfig (line 65) | class TensorDictConfig(metaclass=_TensorDictConfigMeta): function union_tensor_dict (line 72) | def union_tensor_dict(tensor_dict1: TensorDict, tensor_dict2: TensorDict... function union_numpy_dict (line 88) | def union_numpy_dict(tensor_dict1: dict[str, np.ndarray], tensor_dict2: ... function list_of_dict_to_dict_of_list (line 102) | def list_of_dict_to_dict_of_list(list_of_dict: list[dict]): function all_gather_data_proto (line 114) | def all_gather_data_proto(data: TensorDict, process_group): function select_idxs (line 125) | def select_idxs(batch: TensorDict, idxs): FILE: siirl/data_coordinator/sample.py class SampleInfo (line 17) | class SampleInfo(BaseModel): class Sample (line 27) | class Sample(BaseModel): class Config (line 103) | class Config: class SampleManager (line 107) | class SampleManager(BaseModel): class Config (line 111) | class Config: function preprocess_dataloader (line 116) | def preprocess_dataloader(data:Dict, n:int = 1): function Dict2Samples (line 144) | def Dict2Samples(data:TensorDict)-> List[SampleManager]: function Samples2Dict (line 184) | def Samples2Dict(samples: List[Sample]) -> TensorDict: function filter_tensordict (line 251) | def filter_tensordict(batch: TensorDict, indices: List[int]) -> TensorDict: FILE: siirl/engine/actor/base.py class BasePPOActor (line 28) | class BasePPOActor(ABC): method __init__ (line 29) | def __init__(self, config): method compute_log_prob (line 40) | def compute_log_prob(self, data: TensorDict) -> torch.Tensor: method update_policy (line 55) | def update_policy(self, data: TensorDict) -> Dict: FILE: siirl/engine/actor/dp_actor.py class DataParallelPPOActor (line 58) | class DataParallelPPOActor(BasePPOActor): method __init__ (line 59) | def __init__( method _forward_micro_batch (line 92) | def _forward_micro_batch( method _optimizer_step (line 270) | def _optimizer_step(self): method compute_log_prob (line 291) | def compute_log_prob(self, data: TensorDict, calculate_entropy=False) ... method update_policy (line 354) | def update_policy(self, data: TensorDict): FILE: siirl/engine/actor/embodied_actor.py class RobDataParallelPPOActor (line 40) | class RobDataParallelPPOActor(BasePPOActor): method __init__ (line 42) | def __init__( method process_tensor (line 57) | def process_tensor(self, tensor, pad_id): method generate_traj_mask (line 65) | def generate_traj_mask(self, end_step, traj_len): method apply_mask_with_grad_control (line 78) | def apply_mask_with_grad_control(self, log_probs, entropy, mask): method _forward_micro_batch (line 104) | def _forward_micro_batch(self, micro_batch, temperature) -> Tuple[torc... method _forward_micro_batch_update (line 192) | def _forward_micro_batch_update(self, input_ids, attention_mask, pixel... method _forward_micro_batch_entropy (line 248) | def _forward_micro_batch_entropy(self, micro_batch, temperature) -> Tu... method _optimizer_step (line 318) | def _optimizer_step(self): method compute_log_prob (line 328) | def compute_log_prob(self, data: TensorDict, calculate_entropy=False) ... method update_policy (line 380) | def update_policy(self, data: TensorDict): method compute_entropy (line 504) | def compute_entropy(self, bacth_data: TensorDict): FILE: siirl/engine/actor/megatron_actor.py class MegatronPPOActor (line 53) | class MegatronPPOActor(BasePPOActor): method __init__ (line 54) | def __init__( method _validate_config (line 138) | def _validate_config(self, config) -> None: method compute_log_prob (line 149) | def compute_log_prob(self, data: TensorDict, calculate_entropy=False) ... method compute_ppo_loss (line 253) | def compute_ppo_loss(self, model_output, data): method forward_backward_batch (line 307) | def forward_backward_batch( method update_policy (line 513) | def update_policy(self, data:TensorDict) -> dict: FILE: siirl/engine/base_worker/base/worker.py class DistRankInfo (line 29) | class DistRankInfo: class DistGlobalInfo (line 37) | class DistGlobalInfo: class WorkerHelper (line 44) | class WorkerHelper: method _get_node_ip (line 45) | def _get_node_ip(self): method _get_free_port (line 62) | def _get_free_port(self): method get_availale_master_addr_port (line 67) | def get_availale_master_addr_port(self): method _get_pid (line 70) | def _get_pid(self): class Worker (line 75) | class Worker(WorkerHelper): method __new__ (line 85) | def __new__(cls, *args, **kwargs): method _configure_before_init (line 103) | def _configure_before_init(self, register_center_name: str, rank: int): method env_keys (line 134) | def env_keys(cls): method __init__ (line 138) | def __init__(self, cuda_visible_devices=None) -> None: method get_fused_worker_by_name (line 176) | def get_fused_worker_by_name(self, worker_name: str): method _setup_env_cuda_visible_devices (line 185) | def _setup_env_cuda_visible_devices(self): method _configure_with_store (line 232) | def _configure_with_store(self, store: Dict): method get_master_addr_port (line 246) | def get_master_addr_port(self): method get_cuda_visible_devices (line 250) | def get_cuda_visible_devices(self): method world_size (line 258) | def world_size(self): method rank (line 263) | def rank(self): method execute_with_func_generator (line 267) | def execute_with_func_generator(self, func, *args, **kwargs): method execute_func_rank_zero (line 281) | def execute_func_rank_zero(self, func, *args, **kwargs): FILE: siirl/engine/base_worker/megatron/npu_mbridge_patch.py function load_weights_patch (line 20) | def load_weights_patch( function _weight_name_mapping_mcore_local_to_global_patch (line 131) | def _weight_name_mapping_mcore_local_to_global_patch( FILE: siirl/engine/base_worker/megatron/worker.py class MegatronWorker (line 20) | class MegatronWorker(Worker): method __init__ (line 21) | def __init__(self, cuda_visible_devices=None) -> None: method get_megatron_global_info (line 24) | def get_megatron_global_info(self): method get_megatron_rank_info (line 34) | def get_megatron_rank_info(self): method _init_hf_config_and_tf_config (line 45) | def _init_hf_config_and_tf_config( FILE: siirl/engine/base_worker/register_center/register_center.py class WorkerGroupRegisterCenter (line 21) | class WorkerGroupRegisterCenter: method __init__ (line 22) | def __init__(self, rank_zero_info): method get_rank_zero_info (line 27) | def get_rank_zero_info(self): method set_worker_info (line 30) | def set_worker_info(self, rank, node_id) -> None: method get_worker_info (line 33) | def get_worker_info(self) -> Dict[int, str]: function create_worker_group_register_center (line 37) | def create_worker_group_register_center(name, info): FILE: siirl/engine/base_worker/resouce_pool.py class ResourcePool (line 24) | class ResourcePool: method __init__ (line 31) | def __init__(self, process_on_nodes=None, max_colocate_count: int = 10... method add_node (line 45) | def add_node(self, process_count): method world_size (line 49) | def world_size(self): method __call__ (line 53) | def __call__(self) -> Any: method store (line 57) | def store(self): method local_world_size_list (line 60) | def local_world_size_list(self) -> List[int]: method local_rank_list (line 65) | def local_rank_list(self) -> List[int]: class ClassWithInitArgs (line 71) | class ClassWithInitArgs: method __init__ (line 78) | def __init__(self, cls, *args, **kwargs) -> None: method __call__ (line 92) | def __call__(self) -> Any: class WorkerGroup (line 97) | class WorkerGroup: method __init__ (line 105) | def __init__(self, resource_pool: ResourcePool, **kwargs) -> None: method _is_worker_alive (line 124) | def _is_worker_alive(self, worker): method world_size (line 129) | def world_size(self): function get_random_string (line 134) | def get_random_string(length: int) -> str: function sort_placement_group_by_node_ip (line 142) | def sort_placement_group_by_node_ip(pgs: List[PlacementGroup]) -> List[P... class RayResourcePool (line 162) | class RayResourcePool(ResourcePool): method __init__ (line 163) | def __init__( method get_placement_groups (line 180) | def get_placement_groups(self, strategy="STRICT_PACK", name=None, devi... function extract_pg_from_exist (line 208) | def extract_pg_from_exist(resource_pools: Dict[str, RayResourcePool], sr... function merge_resource_pool (line 225) | def merge_resource_pool(rp1: RayResourcePool, rp2: RayResourcePool) -> R... class RayClassWithInitArgs (line 239) | class RayClassWithInitArgs(ClassWithInitArgs): method __init__ (line 247) | def __init__(self, cls, *args, **kwargs) -> None: method set_additional_resource (line 253) | def set_additional_resource(self, additional_resource): method update_options (line 261) | def update_options(self, options: Dict): method __call__ (line 269) | def __call__(self, placement_group, placement_group_bundle_idx, use_gp... FILE: siirl/engine/critic/base.py class BasePPOCritic (line 27) | class BasePPOCritic(ABC): method __init__ (line 28) | def __init__(self, config): method compute_values (line 33) | def compute_values(self, data: TensorDict) -> torch.Tensor: method update_critic (line 38) | def update_critic(self, data: TensorDict): FILE: siirl/engine/critic/dp_critic.py class DataParallelPPOCritic (line 41) | class DataParallelPPOCritic(BasePPOCritic): method __init__ (line 42) | def __init__(self, config: CriticArguments, critic_module: nn.Module, ... method _forward_micro_batch (line 52) | def _forward_micro_batch(self, micro_batch): method _optimizer_step (line 111) | def _optimizer_step(self): method compute_values (line 130) | def compute_values(self, data: TensorDict) -> torch.Tensor: method update_critic (line 169) | def update_critic(self, data: TensorDict): FILE: siirl/engine/critic/megatron_critic.py class MegatronPPOCritic (line 41) | class MegatronPPOCritic(BasePPOCritic): method __init__ (line 42) | def __init__( method _validate_config (line 77) | def _validate_config(self, config) -> None: method compute_values (line 88) | def compute_values(self, data: TensorDict) -> TensorDict: method forward_backward_batch (line 135) | def forward_backward_batch(self, data: TensorDict, forward_only=False,... method update_critic (line 264) | def update_critic(self, data: TensorDict): FILE: siirl/engine/fsdp_workers.py function create_device_mesh_from_group (line 76) | def create_device_mesh_from_group( function get_sharding_strategy (line 154) | def get_sharding_strategy(device_mesh): class ActorRolloutRefWorker (line 166) | class ActorRolloutRefWorker(Worker): method __init__ (line 172) | def __init__(self, config: ActorRolloutRefArguments, role: str, proces... method _build_model_optimizer (line 241) | def _build_model_optimizer( method _build_model_optimizer_legacy (line 328) | def _build_model_optimizer_legacy( method _prepare_and_load_model (line 635) | def _prepare_and_load_model( method _get_model_class (line 750) | def _get_model_class(self, model_config, is_embodied: bool): method _register_embodied_model (line 763) | def _register_embodied_model(self, model_path: str, trust_remote_code:... method _setup_openvla_oft_model (line 806) | def _setup_openvla_oft_model(self, model: torch.nn.Module, model_path:... method _apply_model_modifications (line 828) | def _apply_model_modifications( method _setup_fsdp_wrapper (line 901) | def _setup_fsdp_wrapper( method _create_optimizer_and_scheduler (line 1026) | def _create_optimizer_and_scheduler( method _build_rollout (line 1094) | def _build_rollout(self, trust_remote_code=False): method init_model (line 1156) | def init_model(self): method update_actor (line 1273) | def update_actor(self, data: TensorDict): method generate_sequences (line 1313) | def generate_sequences(self, prompts: TensorDict): method compute_log_prob (line 1357) | def compute_log_prob(self, data: TensorDict): method compute_ref_log_prob (line 1410) | def compute_ref_log_prob(self, data: TensorDict): method save_checkpoint (line 1452) | def save_checkpoint(self, local_path, hdfs_path=None, global_step=0, m... method load_checkpoint (line 1493) | def load_checkpoint(self, local_path, hdfs_path=None, del_local_after_... method set_rollout_sharding_manager (line 1505) | def set_rollout_sharding_manager(self, sharding_manager): class CriticWorker (line 1510) | class CriticWorker(Worker): method __init__ (line 1511) | def __init__(self, config: CriticArguments, process_group: ProcessGroup): method _build_critic_model_optimizer (line 1546) | def _build_critic_model_optimizer(self, config): method init_model (line 1717) | def init_model(self): method compute_values (line 1737) | def compute_values(self, data: TensorDict): method update_critic (line 1759) | def update_critic(self, data: TensorDict): method save_checkpoint (line 1794) | def save_checkpoint(self, local_path, hdfs_path=None, global_step=0, m... method load_checkpoint (line 1806) | def load_checkpoint(self, local_path, hdfs_path=None, del_local_after_... class RewardModelWorker (line 1823) | class RewardModelWorker(Worker): method __init__ (line 1828) | def __init__(self, config: RewardModelArguments, process_group: Proces... method _build_model (line 1855) | def _build_model(self, config: RewardModelArguments): method init_model (line 1935) | def init_model(self): method _forward_micro_batch (line 1940) | def _forward_micro_batch(self, micro_batch): method _expand_to_token_level (line 1986) | def _expand_to_token_level(self, data: TensorDict, scores: torch.Tensor): method _switch_chat_template (line 2001) | def _switch_chat_template(self, data: TensorDict): method compute_rm_score (line 2062) | def compute_rm_score(self, data: TensorDict): class AsyncActorRolloutRefWorker (line 2123) | class AsyncActorRolloutRefWorker(ActorRolloutRefWorker): method _build_rollout (line 2124) | def _build_rollout(self, trust_remote_code=False): method generate_sequences (line 2139) | def generate_sequences(self, prompts: TensorDict): method execute_method (line 2142) | def execute_method(self, method: Union[str, bytes], *args, **kwargs): method chat_completion (line 2148) | async def chat_completion(self, json_request): method wake_up (line 2152) | async def wake_up(self): method sleep (line 2157) | async def sleep(self): method set_rollout_sharding_manager (line 2162) | def set_rollout_sharding_manager(self, sharding_manager): method get_zeromq_address (line 2166) | def get_zeromq_address(self): FILE: siirl/engine/megatron_workers.py function set_random_seed (line 67) | def set_random_seed(seed): class ActorRolloutRefWorker (line 87) | class ActorRolloutRefWorker(MegatronWorker): method __init__ (line 93) | def __init__(self, config: DictConfig, role: str, process_group=None): method _build_model_optimizer (line 161) | def _build_model_optimizer(self, model_path, optim_config, override_mo... method _build_rollout (line 228) | def _build_rollout(self, trust_remote_code=False): method init_model (line 301) | def init_model(self): method update_actor (line 393) | def update_actor(self, data: TensorDict): method generate_sequences (line 427) | def generate_sequences(self, prompts: TensorDict): method load_checkpoint (line 463) | def load_checkpoint(self, checkpoint_path, hdfs_path=None, del_local_a... method load_pretrained_model (line 472) | def load_pretrained_model(self, checkpoint_path, del_local_after_load=... method save_checkpoint (line 475) | def save_checkpoint(self, checkpoint_path, hdfs_path=None, global_step... class AsyncActorRolloutRefWorker (line 484) | class AsyncActorRolloutRefWorker(ActorRolloutRefWorker): method _build_rollout (line 485) | def _build_rollout(self, trust_remote_code=False): method execute_method (line 500) | def execute_method(self, method: Union[str, bytes], *args, **kwargs): method chat_completion (line 506) | async def chat_completion(self, json_request): method wake_up (line 510) | async def wake_up(self): method sleep (line 515) | async def sleep(self): class CriticWorker (line 521) | class CriticWorker(MegatronWorker): method __init__ (line 522) | def __init__(self, config, process_group=None): method _build_critic_model_optimizer (line 569) | def _build_critic_model_optimizer(self, model_path, optim_config, over... method init_model (line 633) | def init_model(self): method compute_values (line 690) | def compute_values(self, data: TensorDict): method update_critic (line 705) | def update_critic(self, data: TensorDict): method load_checkpoint (line 729) | def load_checkpoint(self, local_path, hdfs_path=None, del_local_after_... method save_checkpoint (line 738) | def save_checkpoint(self, local_path, hdfs_path=None, global_step=0, m... class RewardModelWorker (line 746) | class RewardModelWorker(MegatronWorker): method __init__ (line 751) | def __init__(self, config, process_group=None): method _build_rm_model (line 788) | def _build_rm_model(self, model_path, tokenizer, override_model_config... method init_model (line 830) | def init_model(self): method compute_rm_score (line 871) | def compute_rm_score(self, data: TensorDict): function global_initialize_model_parallel (line 885) | def global_initialize_model_parallel(config: ActorRolloutRefArguments): function global_mindspeed_repatch (line 923) | def global_mindspeed_repatch(config): class ActorWorker (line 935) | class ActorWorker(MegatronWorker): method __init__ (line 940) | def __init__(self, config: DictConfig, process_group=None): method _build_actor_model_optimizer (line 959) | def _build_actor_model_optimizer(self, model_path, optim_config, overr... method init_model (line 1018) | def init_model(self): method update_actor (line 1086) | def update_actor(self, data: TensorDict): method compute_log_prob (line 1122) | def compute_log_prob(self, data: TensorDict): method load_checkpoint (line 1156) | def load_checkpoint(self, local_path, hdfs_path=None, del_local_after_... method load_pretrained_model (line 1165) | def load_pretrained_model(self, checkpoint_path, del_local_after_load=... method save_checkpoint (line 1168) | def save_checkpoint(self, local_path, hdfs_path=None, global_step=0, m... class RolloutWorker (line 1177) | class RolloutWorker(MegatronWorker): method __init__ (line 1182) | def __init__(self, config: DictConfig, process_group=None): method _build_rollout (line 1199) | def _build_rollout(self, trust_remote_code=False): method init_model (line 1250) | def init_model(self): method generate_sequences (line 1284) | def generate_sequences(self, prompts: TensorDict): method set_rollout_sharding_manager (line 1304) | def set_rollout_sharding_manager(self, sharding_manager): class ReferenceWorker (line 1308) | class ReferenceWorker(MegatronWorker): method __init__ (line 1313) | def __init__(self, config: DictConfig, process_group=None): method _build_ref_model (line 1331) | def _build_ref_model(self, model_path, override_model_config, override... method init_model (line 1373) | def init_model(self): method compute_ref_log_prob (line 1409) | def compute_ref_log_prob(self, data: TensorDict): class AsyncRolloutWorker (line 1439) | class AsyncRolloutWorker(RolloutWorker): method _build_rollout (line 1440) | def _build_rollout(self, trust_remote_code=False): method execute_method (line 1455) | def execute_method(self, method: Union[str, bytes], *args, **kwargs): method chat_completion (line 1461) | async def chat_completion(self, json_request): method wake_up (line 1465) | async def wake_up(self): method sleep (line 1470) | async def sleep(self): method set_rollout_sharding_manager (line 1475) | def set_rollout_sharding_manager(self, sharding_manager): FILE: siirl/engine/reward_manager/__init__.py function __getattr__ (line 20) | def __getattr__(name): FILE: siirl/engine/reward_manager/dapo.py class DAPORewardManager (line 25) | class DAPORewardManager: method __init__ (line 28) | def __init__( method __call__ (line 48) | def __call__(self, data: TensorDict, return_dict: bool = False): FILE: siirl/engine/reward_manager/embodied.py class EmbodiedRewardManager (line 25) | class EmbodiedRewardManager: method __init__ (line 34) | def __init__( method __call__ (line 79) | def __call__(self, data: TensorDict, return_dict: bool = False) -> Uni... method verify (line 165) | def verify(self, data: TensorDict) -> Tuple[List[float], Dict[str, flo... FILE: siirl/engine/reward_manager/naive.py class NaiveRewardManager (line 25) | class NaiveRewardManager: method __init__ (line 28) | def __init__(self, tokenizer, num_examine, compute_score=None, reward_... method __call__ (line 35) | def __call__(self, data: TensorDict, return_dict=False): FILE: siirl/engine/reward_manager/parallel.py class ParallelRewardManager (line 26) | class ParallelRewardManager: method __init__ (line 29) | def __init__(self, tokenizer, num_examine, compute_score=None, reward_... method _process_single_item (line 35) | def _process_single_item(self, data_item): method _compute_score (line 47) | def _compute_score(self, item): method verify (line 56) | def verify(self, data): method __call__ (line 63) | def __call__(self, data: TensorDict, return_dict: bool = True): FILE: siirl/engine/reward_model/base.py class BasePPORewardModel (line 23) | class BasePPORewardModel(ABC): method __init__ (line 24) | def __init__(self, config): method compute_reward (line 28) | def compute_reward(self, data: TensorDict) -> TensorDict: FILE: siirl/engine/reward_model/megatron/reward_model.py class MegatronRewardModel (line 33) | class MegatronRewardModel(BasePPORewardModel): method __init__ (line 34) | def __init__( method re_encode_by_rm_tokenizer (line 59) | def re_encode_by_rm_tokenizer(self, data: TensorDict) -> TensorDict: method compute_reward (line 124) | def compute_reward(self, data: TensorDict) -> TensorDict: method forward_batch (line 203) | def forward_batch(self, data: TensorDict, use_dynamic_bsz=False, micro... method offload_params_to_cpu (line 304) | def offload_params_to_cpu(self): method load_params_to_cuda (line 312) | def load_params_to_cuda(self): FILE: siirl/engine/rollout/async_server.py function _get_free_port (line 28) | def _get_free_port(): class AsyncServerBase (line 34) | class AsyncServerBase(ABC): method __init__ (line 37) | def __init__(self): method chat_completion (line 42) | async def chat_completion(self, raw_request: Request): method generate (line 50) | async def generate(self, prompt_ids: List[int], sampling_params: Dict[... method init_engine (line 64) | async def init_engine(self): method wake_up (line 69) | async def wake_up(self): method sleep (line 74) | async def sleep(self): function async_server_class (line 79) | def async_server_class( FILE: siirl/engine/rollout/base.py class BaseRollout (line 22) | class BaseRollout(ABC): method generate_sequences (line 26) | def generate_sequences(self, prompts: TensorDict) -> TensorDict: FILE: siirl/engine/rollout/embodied_rollout.py function _timer (line 57) | def _timer(name: str, timing_dict: dict): function crop_and_resize (line 65) | def crop_and_resize(image, crop_scale, batch_size): function center_crop_image (line 111) | def center_crop_image(image): class EmbodiedHFRollout (line 135) | class EmbodiedHFRollout(BaseRollout): method __init__ (line 137) | def __init__(self, module: nn.Module, config: ActorRolloutRefArguments): method close (line 177) | def close(self): method __del__ (line 184) | def __del__(self): method embodied_preprocess (line 188) | def embodied_preprocess(self): method generate_sequences (line 202) | def generate_sequences(self, prompts): method process_input (line 247) | def process_input(self,inputs:list, task_descriptions:list): method _generate_chunk_rollout (line 312) | def _generate_chunk_rollout(self, prompts): method _generate_one_step (line 498) | def _generate_one_step(self, prompts: dict): method _obs_to_input (line 660) | def _obs_to_input(self, obs): FILE: siirl/engine/rollout/hf_rollout.py class HFRollout (line 38) | class HFRollout(BaseRollout): method __init__ (line 39) | def __init__(self, module: nn.Module, config): method generate_sequences (line 44) | def generate_sequences(self, prompts: TensorDict) -> TensorDict: method _generate_minibatch (line 53) | def _generate_minibatch(self, prompts: TensorDict) -> TensorDict: FILE: siirl/engine/rollout/schemas.py class FinishReasonTypeEnum (line 36) | class FinishReasonTypeEnum(str, Enum): method from_str (line 44) | def from_str(cls, value: str) -> "FinishReasonTypeEnum": class Message (line 55) | class Message(BaseModel): class AsyncRolloutRequestStateEnum (line 61) | class AsyncRolloutRequestStateEnum(str, Enum): class TokenizationSanityCheckModeEnum (line 72) | class TokenizationSanityCheckModeEnum(str, Enum): class AsyncRolloutRequest (line 80) | class AsyncRolloutRequest(BaseModel): method initialize_request (line 122) | def initialize_request(cls, values): method _handle_apply_chat_template (line 219) | def _handle_apply_chat_template( method _get_position_ids (line 255) | def _get_position_ids( method _update_input_ids (line 293) | def _update_input_ids( method _update_multi_modal_inputs (line 330) | def _update_multi_modal_inputs(self, new_multi_modal_inputs: Dict[str,... method get_generation_prompt_ids (line 342) | def get_generation_prompt_ids( method add_user_message (line 373) | def add_user_message( method add_assistant_message (line 389) | def add_assistant_message( method add_tool_response_messages (line 407) | def add_tool_response_messages( method update_metrics (line 484) | def update_metrics(self, metrics: Any, tool_id: str) -> None: method _get_prompt_diffs (line 492) | def _get_prompt_diffs( method finalize (line 549) | def finalize( method truncate_output_ids (line 660) | def truncate_output_ids( FILE: siirl/engine/rollout/sglang_rollout/async_sglang_server.py class AsyncSglangServer (line 32) | class AsyncSglangServer(AsyncServerBase): method __init__ (line 33) | def __init__(self, config: ActorRolloutRefArguments, spmd_engine: SGLa... method init_engine (line 41) | async def init_engine(self): method chat_completion (line 48) | async def chat_completion(self, raw_request: Request): method generate (line 57) | async def generate(self, prompt_ids: List[int], sampling_params: Dict[... method wake_up (line 60) | async def wake_up(self): method sleep (line 70) | async def sleep(self): FILE: siirl/engine/rollout/sglang_rollout/sglang_rollout.py function _set_envs_and_config (line 93) | def _set_envs_and_config(server_args: ServerArgs): class AsyncEngine (line 137) | class AsyncEngine(sglang.srt.entrypoints.engine.Engine): method __init__ (line 138) | def __init__(self, **kwargs): method release_memory_occupation (line 143) | async def release_memory_occupation(self, tags: Optional[list[str]] = ... method resume_memory_occupation (line 151) | async def resume_memory_occupation(self, tags: Optional[list[str]] = N... method update_weights_from_tensor (line 166) | async def update_weights_from_tensor( method flush_cache (line 183) | async def flush_cache(self): function _pre_process_inputs (line 189) | def _pre_process_inputs( function _post_process_outputs (line 199) | def _post_process_outputs(processing_class, output): function get_tool_call_parser_type (line 228) | def get_tool_call_parser_type( class SGLangRollout (line 252) | class SGLangRollout(BaseRollout): method __init__ (line 253) | def __init__( method _init_distributed_env (line 329) | def _init_distributed_env(self, device_mesh_cpu, **kwargs): method _verify_config (line 374) | def _verify_config(self, model_hf_config): method _init_inference_engine (line 420) | def _init_inference_engine(self, trust_remote_code, actor_module, port): method _init_sampling_params (line 479) | def _init_sampling_params(self, **kwargs): method _initialize_tools (line 495) | def _initialize_tools(self, config, processing_class): method _initialize_interactions (line 541) | def _initialize_interactions(self, config): method generate_sequences (line 558) | def generate_sequences(self, prompts: TensorDict, **kwargs) -> TensorD... method _batch_level_generate_sequences (line 585) | def _batch_level_generate_sequences(self, prompts: TensorDict, **kwarg... method _async_rollout_a_request (line 775) | async def _async_rollout_a_request( method _handle_engine_call (line 982) | async def _handle_engine_call( method _handle_engine_generate (line 988) | async def _handle_engine_generate( method _handle_pending_state (line 1003) | async def _handle_pending_state(self, _req: AsyncRolloutRequest) -> As... method generate_sequences_with_tools (line 1026) | def generate_sequences_with_tools(self, prompts: TensorDict, **kwargs)... method _req_level_generate_sequences (line 1036) | def _req_level_generate_sequences(self, prompts: TensorDict, **kwargs)... method _preprocess_prompt_to_async_rollout_requests (line 1208) | def _preprocess_prompt_to_async_rollout_requests(self, prompts: Tensor... method chat_completion (line 1277) | async def chat_completion(self, json_request): method generate (line 1352) | async def generate( method wake_up (line 1360) | async def wake_up(self): method sleep (line 1367) | async def sleep(self): method _init_zeromq (line 1375) | def _init_zeromq(self) -> str: method _get_free_port (line 1398) | def _get_free_port(self): method _loop_forever (line 1405) | def _loop_forever(self): method get_zeromq_address (line 1411) | def get_zeromq_address(self): method execute_method (line 1414) | def execute_method(self, method: Union[str, bytes], *args, **kwargs): method get_device_mesh (line 1426) | def get_device_mesh(self): function ensure_event_loop (line 1429) | def ensure_event_loop(): FILE: siirl/engine/rollout/sglang_rollout/utils.py function broadcast_pyobj (line 26) | def broadcast_pyobj( FILE: siirl/engine/rollout/vllm_rollout/__init__.py function get_version (line 20) | def get_version(pkg): FILE: siirl/engine/rollout/vllm_rollout/vllm_async_server.py class ExternalZeroMQDistributedExecutor (line 41) | class ExternalZeroMQDistributedExecutor(Executor): method _init_executor (line 46) | def _init_executor(self) -> None: method collective_rpc (line 67) | def collective_rpc( method check_health (line 89) | def check_health(self): class AsyncvLLMServer (line 92) | class AsyncvLLMServer(AsyncServerBase): method __init__ (line 108) | def __init__(self, config: ActorRolloutRefArguments, spmd_engine: Any... method init_engine (line 122) | def init_engine(self): method _create_engine_config (line 179) | def _create_engine_config(self, engine_args: AsyncEngineArgs, zmq_addr... method generate_sequences (line 192) | def generate_sequences(self, prompts: TensorDict, **kwargs) -> TensorD... method chat_completion (line 195) | async def chat_completion(self, raw_request: Request): method generate (line 212) | async def generate(self, prompt_ids: List[int], sampling_params: Dict[... method wake_up (line 226) | async def wake_up(self): method sleep (line 230) | async def sleep(self): FILE: siirl/engine/rollout/vllm_rollout/vllm_rollout_spmd.py function _pre_process_inputs (line 83) | def _pre_process_inputs(pad_token_id, prompt_token_ids: torch.Tensor) ->... function _repeat_interleave (line 92) | def _repeat_interleave(value: Union[torch.Tensor, np.ndarray], repeats: ... class vLLMRollout (line 99) | class vLLMRollout(BaseRollout): method __init__ (line 100) | def __init__(self, model_path: str, config: RolloutArguments, tokenize... method update_sampling_params (line 248) | def update_sampling_params(self, **kwargs): method generate_sequences (line 265) | def generate_sequences(self, prompts: TensorDict, **kwargs) -> TensorD... function _monkey_patch_compute_logits (line 426) | def _monkey_patch_compute_logits(model, vocab_size: int): class vLLMAsyncRollout (line 440) | class vLLMAsyncRollout: method __init__ (line 444) | def __init__(self, model_path: str, config: DictConfig, tokenizer, mod... method _init_zeromq (line 454) | def _init_zeromq(self) -> str: method _get_free_port (line 477) | def _get_free_port(self): method _loop_forever (line 484) | def _loop_forever(self): method get_zeromq_address (line 491) | def get_zeromq_address(self): method init_worker (line 494) | def init_worker(self, all_kwargs: List[Dict[str, Any]]): method load_model (line 504) | def load_model(self, *args, **kwargs): method sleep (line 512) | def sleep(self, *args, **kwargs): method wake_up (line 519) | def wake_up(self, *args, **kwargs): method execute_method (line 526) | def execute_method(self, method: Union[str, bytes], *args, **kwargs): FILE: siirl/engine/sharding_manager/base.py class BaseShardingManager (line 21) | class BaseShardingManager: method __enter__ (line 22) | def __enter__(self): method __exit__ (line 25) | def __exit__(self, exc_type, exc_value, traceback): method preprocess_data (line 28) | def preprocess_data(self, data: TensorDict) -> TensorDict: method postprocess_data (line 31) | def postprocess_data(self, data: TensorDict) -> TensorDict: FILE: siirl/engine/sharding_manager/fsdp_hf.py class FSDPHFShardingManager (line 28) | class FSDPHFShardingManager(BaseShardingManager): method __init__ (line 40) | def __init__( method __enter__ (line 69) | def __enter__(self): method __exit__ (line 88) | def __exit__(self, exc_type, exc_value, traceback): FILE: siirl/engine/sharding_manager/fsdp_sglang.py function _preprocess_tensor_for_update_weights (line 58) | def _preprocess_tensor_for_update_weights(tensor: torch.Tensor): class MultiAgentFSDPSGLangShardingManager (line 64) | class MultiAgentFSDPSGLangShardingManager(BaseShardingManager): method __init__ (line 66) | def __init__( method __enter__ (line 108) | def __enter__(self): method __exit__ (line 133) | def __exit__(self, exc_type, exc_value, traceback): method update_weights (line 149) | async def update_weights(self, params): method release_memory (line 182) | async def release_memory(self): method wake_up (line 186) | async def wake_up(self): method sleep (line 210) | async def sleep(self): method preprocess_data (line 225) | def preprocess_data(self, data: TensorDict) -> TensorDict: method postprocess_data (line 236) | def postprocess_data(self, data: TensorDict) -> TensorDict: FILE: siirl/engine/sharding_manager/fsdp_ulysses.py class FSDPUlyssesShardingManager (line 26) | class FSDPUlyssesShardingManager(BaseShardingManager): method __init__ (line 31) | def __init__(self, device_mesh: DeviceMesh): method __enter__ (line 36) | def __enter__(self): method __exit__ (line 44) | def __exit__(self, exc_type, exc_value, traceback): method preprocess_data (line 51) | def preprocess_data(self, data: TensorDict) -> TensorDict: method postprocess_data (line 63) | def postprocess_data(self, data: TensorDict) -> TensorDict: FILE: siirl/engine/sharding_manager/fsdp_vllm.py class MultiAgentFSDPVLLMShardingManager (line 42) | class MultiAgentFSDPVLLMShardingManager(BaseShardingManager): method __init__ (line 44) | def __init__(self, module: FSDP, inference_engine: LLM, model_config, ... method __enter__ (line 86) | def __enter__(self): method __exit__ (line 180) | def __exit__(self, exc_type, exc_value, traceback): method update_params (line 193) | def update_params(self, updated_params, peft_config=None): FILE: siirl/engine/sharding_manager/megatron_sglang.py function _preprocess_tensor_for_update_weights (line 54) | def _preprocess_tensor_for_update_weights(tensor: torch.Tensor): class MultiAgentMegatronSGLangShardingManager (line 59) | class MultiAgentMegatronSGLangShardingManager(BaseShardingManager): method __init__ (line 60) | def __init__( method __enter__ (line 101) | def __enter__(self): method __exit__ (line 106) | def __exit__(self, exc_type, exc_value, traceback): method update_weights (line 110) | async def update_weights(self, params): method release_memory (line 157) | async def release_memory(self): method wake_up (line 162) | async def wake_up(self): method sleep (line 190) | async def sleep(self): FILE: siirl/engine/sharding_manager/megatron_vllm.py class AllGatherPPModel (line 55) | class AllGatherPPModel: method __init__ (line 56) | def __init__(self, model_provider, use_distributed_optimizer=True) -> ... method _build_param_buffer (line 105) | def _build_param_buffer(self, pp_rank): method _build_param_references (line 121) | def _build_param_references(self, pp_rank, maintain_weight=False): method _load_params_to_cuda (line 127) | def _load_params_to_cuda(self, pp_rank, to_empty=False): method _offload_params_to_cpu (line 137) | def _offload_params_to_cpu(self, pp_rank, to_empty=False): method load_params_to_cuda (line 147) | def load_params_to_cuda(self, to_empty=False): method allgather_params (line 153) | def allgather_params(self): method forward (line 163) | def forward(self, *inputs, **kwargs): method __call__ (line 182) | def __call__(self, *inputs, **kwargs): method eval (line 185) | def eval(self): method train (line 189) | def train(self): method offload_params_to_cpu (line 193) | def offload_params_to_cpu(self, to_empty=False): method get_all_params (line 199) | def get_all_params(self): method update_this_rank_models (line 222) | def update_this_rank_models(self, new_models): method this_rank_models (line 227) | def this_rank_models(self): method pp_size (line 231) | def pp_size(self): method pp_rank (line 235) | def pp_rank(self): method pp_group (line 239) | def pp_group(self): method pp_models (line 243) | def pp_models(self): class MultiAgentMegatronVLLMShardingManager (line 265) | class MultiAgentMegatronVLLMShardingManager(BaseShardingManager): method __init__ (line 267) | def __init__( method __enter__ (line 322) | def __enter__(self): method __exit__ (line 376) | def __exit__(self, exc_type, exc_value, traceback): FILE: siirl/environment/embodied/adapters/libero.py class LIBEROAdapter (line 38) | class LIBEROAdapter(BaseVLAEnvironment): method __init__ (line 51) | def __init__(self, method _blocking_reset (line 87) | def _blocking_reset(self, task_ids: Optional[List[int]] = None, trial_... method reset (line 205) | async def reset(self, task_ids: Optional[List[int]] = None, trial_ids:... method _blocking_step (line 209) | def _blocking_step(self, action: Dict[str, Any]) -> List[Dict[str, Any]]: method step (line 284) | async def step(self, action: Dict[str, Any]) -> List[Dict[str, Any]]: method close (line 291) | def close(self): method _get_libero_env (line 298) | def _get_libero_env(task, gpu_id, resolution=256): method _get_dummy_action (line 313) | def _get_dummy_action(self) -> List[float]: method _normalize_gripper_action (line 317) | def _normalize_gripper_action(self, action: np.ndarray, binarize: bool... method _invert_gripper_action (line 342) | def _invert_gripper_action(self, action: np.ndarray) -> np.ndarray: FILE: siirl/environment/embodied/base.py class BaseVLAEnvironment (line 24) | class BaseVLAEnvironment(ABC): method reset (line 32) | async def reset(self) -> Dict[str, Any]: method step (line 43) | async def step(self, action: Dict[str, Any]) -> Tuple[Dict[str, Any], ... FILE: siirl/environment/embodied/venv.py function deprecation (line 44) | def deprecation(msg: str) -> None: class CloudpickleWrapper (line 49) | class CloudpickleWrapper(object): method __init__ (line 52) | def __init__(self, data: Any) -> None: method __getstate__ (line 55) | def __getstate__(self) -> str: method __setstate__ (line 58) | def __setstate__(self, data: str) -> None: class EnvWorker (line 78) | class EnvWorker(ABC): method __init__ (line 81) | def __init__(self, env_fn: Callable[[], gym.Env]) -> None: method get_env_attr (line 94) | def get_env_attr(self, key: str) -> Any: method set_env_attr (line 98) | def set_env_attr(self, key: str, value: Any) -> None: method send (line 101) | def send(self, action: Optional[np.ndarray]) -> None: method recv (line 120) | def recv( method reset (line 145) | def reset(self, **kwargs: Any) -> Union[np.ndarray, Tuple[np.ndarray, ... method step (line 148) | def step( method wait (line 161) | def wait( method seed (line 167) | def seed(self, seed: Optional[int] = None) -> Optional[List[int]]: method render (line 172) | def render(self, **kwargs: Any) -> Any: method close_env (line 177) | def close_env(self) -> None: method close (line 180) | def close(self) -> None: class ShArray (line 187) | class ShArray: method __init__ (line 190) | def __init__(self, dtype: np.generic, shape: Tuple[int]) -> None: method save (line 195) | def save(self, ndarray: np.ndarray) -> None: method get (line 203) | def get(self) -> np.ndarray: function _setup_buf (line 208) | def _setup_buf(space: gym.Space) -> Union[dict, tuple, ShArray]: function _worker (line 219) | def _worker( class DummyEnvWorker (line 307) | class DummyEnvWorker(EnvWorker): method __init__ (line 310) | def __init__(self, env_fn: Callable[[], gym.Env]) -> None: method get_env_attr (line 314) | def get_env_attr(self, key: str) -> Any: method set_env_attr (line 317) | def set_env_attr(self, key: str, value: Any) -> None: method reset (line 320) | def reset(self, **kwargs: Any) -> Union[np.ndarray, Tuple[np.ndarray, ... method wait (line 326) | def wait( # type: ignore method send (line 332) | def send(self, action: Optional[np.ndarray], **kwargs: Any) -> None: method seed (line 338) | def seed(self, seed: Optional[int] = None) -> Optional[List[int]]: method render (line 346) | def render(self, **kwargs: Any) -> Any: method close_env (line 349) | def close_env(self) -> None: method check_success (line 352) | def check_success(self): method get_segmentation_of_interest (line 355) | def get_segmentation_of_interest(self, segmentation_image): method get_sim_state (line 358) | def get_sim_state(self): method set_init_state (line 361) | def set_init_state(self, init_state): class SubprocEnvWorker (line 365) | class SubprocEnvWorker(EnvWorker): method __init__ (line 368) | def __init__( method get_env_attr (line 391) | def get_env_attr(self, key: str) -> Any: method set_env_attr (line 395) | def set_env_attr(self, key: str, value: Any) -> None: method _decode_obs (line 398) | def _decode_obs(self) -> Union[dict, tuple, np.ndarray]: method wait (line 414) | def wait( # type: ignore method send_reinit (line 433) | def send_reinit(self, env_fn: Callable[[], gym.Env]) -> None: method recv_reinit (line 440) | def recv_reinit(self) -> bool: method send (line 449) | def send(self, action: Optional[np.ndarray], **kwargs: Any) -> None: method recv (line 457) | def recv( method reset (line 482) | def reset(self, **kwargs: Any) -> Union[np.ndarray, Tuple[np.ndarray, ... method seed (line 499) | def seed(self, seed: Optional[int] = None) -> Optional[List[int]]: method render (line 505) | def render(self, **kwargs: Any) -> Any: method close_env (line 509) | def close_env(self) -> None: method check_success (line 520) | def check_success(self): method get_segmentation_of_interest (line 524) | def get_segmentation_of_interest(self, segmentation_image): method get_sim_state (line 528) | def get_sim_state(self): method set_init_state (line 532) | def set_init_state(self, init_state): class BaseVectorEnv (line 547) | class BaseVectorEnv(object): method __init__ (line 598) | def __init__( method _assert_is_not_closed (line 633) | def _assert_is_not_closed(self) -> None: method __len__ (line 638) | def __len__(self) -> int: method __getattribute__ (line 642) | def __getattribute__(self, key: str) -> Any: method get_env_attr (line 654) | def get_env_attr( method set_env_attr (line 678) | def set_env_attr( method _wrap_id (line 701) | def _wrap_id( method _assert_id (line 709) | def _assert_id(self, id: Union[List[int], np.ndarray]) -> None: method reset (line 718) | def reset( method step (line 765) | def step( method seed (line 859) | def seed( method render (line 882) | def render(self, **kwargs: Any) -> List[Any]: method close (line 892) | def close(self) -> None: class DummyVectorEnv (line 904) | class DummyVectorEnv(BaseVectorEnv): method __init__ (line 912) | def __init__(self, env_fns: List[Callable[[], gym.Env]], **kwargs: Any... method check_success (line 915) | def check_success(self): method get_segmentation_of_interest (line 918) | def get_segmentation_of_interest(self, segmentation_images): method get_sim_state (line 924) | def get_sim_state(self): method set_init_state (line 927) | def set_init_state( class SubprocVectorEnv (line 952) | class SubprocVectorEnv(BaseVectorEnv): method __init__ (line 960) | def __init__(self, env_fns: List[Callable[[], gym.Env]], **kwargs: Any... method reinit_envs (line 966) | def reinit_envs(self, env_fns: List[Callable[[], gym.Env]], id: Option... method check_success (line 993) | def check_success(self): method get_segmentation_of_interest (line 996) | def get_segmentation_of_interest(self, segmentation_images): method get_sim_state (line 1002) | def get_sim_state(self): method set_init_state (line 1005) | def set_init_state( FILE: siirl/execution/dag/builtin_pipelines.py function grpo_pipeline (line 27) | def grpo_pipeline() -> TaskGraph: function ppo_pipeline (line 87) | def ppo_pipeline() -> TaskGraph: function dapo_pipeline (line 161) | def dapo_pipeline() -> TaskGraph: function embodied_srpo_pipeline (line 230) | def embodied_srpo_pipeline() -> TaskGraph: FILE: siirl/execution/dag/config_loader.py class Ref (line 25) | class Ref: method __init__ (line 30) | def __init__(self, path): method __repr__ (line 33) | def __repr__(self): function ref_constructor (line 37) | def ref_constructor(loader, node): function resolve_refs (line 54) | def resolve_refs(config_item: Any, global_config: Dict[str, Any]) -> Any: class DAGConfigLoader (line 92) | class DAGConfigLoader: method __init__ (line 97) | def __init__(self): method _parse_raw_config (line 101) | def _parse_raw_config(raw_dag_config: Dict[str, Any], file_path: str) ... method load_from_file (line 206) | def load_from_file(file_path: str, file_type: str = "yaml") -> TaskGraph: FILE: siirl/execution/dag/node.py function dynamic_load_function (line 28) | def dynamic_load_function(func_path: str): class NodeType (line 44) | class NodeType(Enum): class NodeRole (line 60) | class NodeRole(Enum): class NodeStatus (line 77) | class NodeStatus(Enum): class AgentProcess (line 90) | class AgentProcess: method __init__ (line 91) | def __init__(self, agent_options: AgentArguments, node_config): method load_attr (line 119) | def load_attr(self, file_path, attr_name): method init_env_class (line 134) | def init_env_class(self): method _init_process_handle (line 141) | def _init_process_handle(self, process_path): method apply_pre_process (line 148) | def apply_pre_process(self, prompt: Optional[Tuple[str, List]], obs: O... method apply_post_process (line 176) | def apply_post_process(self, oridinal_prompt, templated_prompt, respon... class Node (line 211) | class Node: method __init__ (line 216) | def __init__( method _resolve_executable (line 289) | def _resolve_executable(self) -> None: method executable (line 323) | def executable(self) -> Optional[Callable]: method executable (line 328) | def executable(self, execute: Optional[Callable]): method add_dependency (line 332) | def add_dependency(self, dependency_id: str) -> None: method remove_dependency (line 341) | def remove_dependency(self, dependency_id: str) -> None: method is_ready (line 350) | def is_ready(self, completed_node_ids: Set[str]) -> bool: method update_status (line 362) | def update_status(self, new_status: NodeStatus, error_info: Optional[s... method update_config (line 372) | def update_config(self, new_config_items: Dict[str, Any], overwrite: b... method can_retry (line 400) | def can_retry(self) -> bool: method increment_retry_count (line 404) | def increment_retry_count(self) -> None: method run (line 408) | def run(self, **kwargs: Any) -> Any: method __repr__ (line 475) | def __repr__(self) -> str: method copy (line 482) | def copy(self) -> "Node": FILE: siirl/execution/dag/pipeline.py class NodeConfig (line 31) | class NodeConfig: class Pipeline (line 37) | class Pipeline: method __init__ (line 59) | def __init__(self, pipeline_id: str, description: str = ""): method add_node (line 71) | def add_node( method build (line 113) | def build(self) -> TaskGraph: method visualize (line 165) | def visualize(self, output_path: str = None, directory: str = "./"): FILE: siirl/execution/dag/task_graph.py class TaskGraph (line 23) | class TaskGraph: method __init__ (line 28) | def __init__(self, graph_id: str): method add_node (line 42) | def add_node(self, node: Node) -> None: method add_nodes (line 59) | def add_nodes(self, nodes: List[Node]) -> None: method _update_adj_for_node (line 70) | def _update_adj_for_node(self, node: Node) -> None: method build_adjacency_lists (line 89) | def build_adjacency_lists(self) -> None: method get_node (line 113) | def get_node(self, node_id: str) -> Optional[Node]: method get_dependencies (line 123) | def get_dependencies(self, node_id: str) -> List[Node]: method get_dependents (line 140) | def get_dependents(self, node_id: str) -> List[Node]: method get_downstream_nodes (line 155) | def get_downstream_nodes(self, node_id: str) -> List[Node]: method get_entry_nodes (line 159) | def get_entry_nodes(self) -> List[Node]: method get_exit_nodes (line 169) | def get_exit_nodes(self) -> List[Node]: method validate_graph (line 179) | def validate_graph(self) -> Tuple[bool, Optional[str]]: method get_topological_sort (line 200) | def get_topological_sort(self) -> List[str]: method reset_nodes_status (line 247) | def reset_nodes_status(self) -> None: method load_from_config (line 256) | def load_from_config(cls, graph_id: str, config_data: List[Dict[str, A... method __repr__ (line 307) | def __repr__(self) -> str: method copy (line 359) | def copy(self) -> "TaskGraph": method save_dag_pic (line 366) | def save_dag_pic(self, filename: str = "task_graph", directory: Option... method get_nodes_by_type (line 440) | def get_nodes_by_type(self, node_types: List[NodeType]) -> List[Node]: method get_nodes_by_role (line 452) | def get_nodes_by_role(self, node_role: NodeRole) -> List[Node]: FILE: siirl/execution/dag/task_loader.py function generate_structural_signature (line 24) | def generate_structural_signature(graph: TaskGraph) -> str: function get_all_downstream_nodes_recursive (line 64) | def get_all_downstream_nodes_recursive(src_task_graph: TaskGraph, start_... function get_all_ancestors (line 93) | def get_all_ancestors(graph: TaskGraph, node_id: str) -> Set[str]: function find_all_paths_dfs (line 129) | def find_all_paths_dfs(src_task_graph: TaskGraph, current_node_id: str, ... function find_all_paths (line 157) | def find_all_paths(src_task_graph: TaskGraph, start_node_id: str, end_no... function split_single_structure (line 180) | def split_single_structure(src_task_graph: TaskGraph, parallel_branch_no... function split_by_fan_out_to_exits (line 263) | def split_by_fan_out_to_exits(src_task_graph: TaskGraph, naming_prefix_i... function split_by_reconverging_paths (line 357) | def split_by_reconverging_paths(src_task_graph: TaskGraph, naming_prefix... function discover_and_split_parallel_paths (line 464) | def discover_and_split_parallel_paths(src_task_graph: TaskGraph) -> List... FILE: siirl/execution/metric_worker/metric_worker.py class MetricClient (line 34) | class MetricClient(): method __init__ (line 40) | def __init__(self, metric_worker: ActorHandle): method stop (line 49) | def stop(self): method submit_metric (line 54) | def submit_metric(self, metrics: dict, world_size): method wait_submit (line 64) | def wait_submit(self): method wait_final_res (line 69) | def wait_final_res(self): method compute_local_data_metric (line 78) | def compute_local_data_metric(self, data: TensorDict, world_size: int): method compute_local_throughout_metrics (line 98) | def compute_local_throughout_metrics(self, data: TensorDict, timing_ra... method compute_local_timing_metrics (line 113) | def compute_local_timing_metrics(self, data: TensorDict, timing_raw: d... method process_local_validation_metrics (line 125) | def process_local_validation_metrics(self, data_sources: list[str], sa... class MetricWorker (line 144) | class MetricWorker: method __init__ (line 150) | def __init__(self) -> None: method start (line 158) | async def start(self): method submit_metric (line 170) | async def submit_metric(self, metric: dict, world_size: int): method stop (line 181) | async def stop(self): method compute_metric (line 191) | async def compute_metric(self, metric_name, metrics): method parse_metric (line 211) | async def parse_metric(self, metric_data: tuple): method _process_metrics_loop (line 236) | async def _process_metrics_loop(self): method wait_final_res (line 246) | async def wait_final_res(self): FILE: siirl/execution/metric_worker/utils.py class Metric (line 24) | class Metric: function MetricFunc (line 29) | def MetricFunc(name: str): function SumMetric (line 39) | def SumMetric(metrics: List[Metric]): function MeanMetric (line 45) | def MeanMetric(metrics: List[Metric]): function MaxMetric (line 52) | def MaxMetric(metrics: List[Metric]): function MinMetric (line 58) | def MinMetric(metrics: List[Metric]): FILE: siirl/execution/rollout_flow/multi_agent/multiagent_generate.py class MultiAgentLoop (line 39) | class MultiAgentLoop(): method __init__ (line 40) | def __init__(self, dag, config: ActorRolloutRefArguments, node_workers... method _parse_graph (line 60) | def _parse_graph(self, graph:TaskGraph): method _generate_node_worker_key (line 78) | def _generate_node_worker_key(self, node: Node) -> str: method node_if_local (line 82) | def node_if_local(self, node): method _preprocess (line 86) | def _preprocess(self, batch:TensorDict) -> List[str]: method _generate_key (line 99) | def _generate_key(self, cur_node: Node, next_node: Node, batch_id: int... method async_put_data (line 127) | async def async_put_data(self, key: str, value: Tuple[AgentOutput, str... method async_get_envdata (line 154) | async def async_get_envdata(self, key: str, timing_raw: Dict[str, floa... method async_get_data (line 177) | async def async_get_data(self, key: str, timing_raw: Dict[str, float]): method spread_task (line 207) | async def spread_task(self, cur_node, node_idx, batch_idx): method generate_spread (line 231) | async def generate_spread(self): method check_colocate_running (line 250) | async def check_colocate_running(self, finished_res: Dict, visited_age... method colocate_task (line 277) | async def colocate_task(self, agent_output:AgentOutput, agent_res:Dict... method generate_colocate (line 396) | async def generate_colocate(self, bs, sampling_params: Dict[str, Any],... method _postprocess (line 486) | def _postprocess(self, agent_outputs: Dict[str, List[AgentOutput]], me... method generate_sequence (line 658) | def generate_sequence(self, batch:TensorDict, timing_raw: Dict[str, fl... FILE: siirl/execution/rollout_flow/multi_agent/utils.py class AgentOutputStatus (line 4) | class AgentOutputStatus: class AgentOutput (line 10) | class AgentOutput(BaseModel): FILE: siirl/execution/rollout_flow/multiturn/agent_loop/agent_loop.py function get_ddp_world_size_rank (line 43) | async def get_ddp_world_size_rank(local_world_size, local_rank, local_pa... class AsyncLLMServerManager (line 51) | class AsyncLLMServerManager: method __init__ (line 58) | def __init__(self, config: DictConfig, server, max_cache_size: int = 1... method generate (line 69) | async def generate( class AgentLoopMetrics (line 95) | class AgentLoopMetrics(BaseModel): class AgentLoopOutput (line 102) | class AgentLoopOutput(BaseModel): class AgentLoopBase (line 112) | class AgentLoopBase(ABC): method __init__ (line 118) | def __init__(self, config: DictConfig, server_manager: AsyncLLMServerM... method init_class (line 133) | def init_class(cls, config: DictConfig, tokenizer: AutoTokenizer): method run (line 140) | async def run(self, messages: List[Dict[str, Any]], sampling_params: D... class AgentLoopWorker (line 153) | class AgentLoopWorker: method __init__ (line 156) | def __init__(self, config: DictConfig, server_handles: List[ray.actor.... method generate_sequences (line 171) | async def generate_sequences(self, batch: TensorDict) -> TensorDict: method _run_agent_loop (line 220) | async def _run_agent_loop( method get_agent_loop_class (line 228) | def get_agent_loop_class(self, agent_name: str) -> Type[AgentLoopBase]: method _postprocess (line 239) | def _postprocess(self, inputs: List[AgentLoopOutput]) -> TensorDict: class AgentLoopManager (line 305) | class AgentLoopManager: method __init__ (line 308) | def __init__(self, config: DictConfig, cur_dp_rank, name_prefix, engin... method _initialize_llm_servers (line 323) | def _initialize_llm_servers(self): method _init_agent_loop_workers (line 339) | def _init_agent_loop_workers(self): method generate_sequences (line 343) | async def generate_sequences(self, prompts: TensorDict) -> TensorDict: method _performance_metrics (line 365) | def _performance_metrics(self, metrics: List[List[Dict[str, str]]], ou... method wake_up (line 387) | async def wake_up(self): method sleep (line 391) | async def sleep(self): FILE: siirl/execution/rollout_flow/multiturn/agent_loop/single_turn_agent_loop.py function _timer (line 27) | def _timer(name: str, timing_dict: dict): class SingleTurnAgentLoop (line 39) | class SingleTurnAgentLoop(AgentLoopBase): method __init__ (line 42) | def __init__(self, config, server_manager, tokenizer): method run (line 47) | async def run(self, messages: List[Dict[str, Any]], sampling_params: D... FILE: siirl/execution/rollout_flow/multiturn/agent_loop/tool_agent_loop.py function _timer (line 34) | def _timer(name: str, timing_dict: dict): class FunctionCall (line 46) | class FunctionCall(BaseModel): class ToolParser (line 59) | class ToolParser(ABC): method extract_tool_calls (line 61) | async def extract_tool_calls(self, responses_ids: List[int], prompt_id... class HermesToolParser (line 73) | class HermesToolParser(ToolParser): method __init__ (line 76) | def __init__(self, tokenizer) -> None: method extract_tool_calls (line 83) | async def extract_tool_calls(self, responses_ids: List[int], prompt_id... class ToolAgentLoop (line 101) | class ToolAgentLoop(AgentLoopBase): method __init__ (line 102) | def __init__(self, config, server_manager, tokenizer): method init_class (line 106) | def init_class(cls, config, tokenizer): method run (line 130) | async def run(self, messages: List[Dict[str, Any]], sampling_params: D... method _call_tool (line 208) | async def _call_tool(self, tool_call: FunctionCall) -> Dict[str, str]: method get_tool_parser (line 241) | def get_tool_parser(cls, name: str) -> ToolParser: FILE: siirl/execution/rollout_flow/multiturn/interactions/base.py class BaseInteraction (line 20) | class BaseInteraction: method __init__ (line 21) | def __init__(self, config: Dict[str, Any]): method start_interaction (line 25) | async def start_interaction(self, instance_id: Optional[str] = None, *... method generate_response (line 39) | async def generate_response( method calculate_score (line 56) | async def calculate_score(self) -> float: # More clear score calculat... method finalize_interaction (line 66) | async def finalize_interaction(self) -> None: # More clear interactio... FILE: siirl/execution/rollout_flow/multiturn/interactions/gsm8k_interaction.py class Gsm8kInteraction (line 28) | class Gsm8kInteraction(BaseInteraction): method __init__ (line 37) | def __init__(self, config: dict): method start_interaction (line 41) | async def start_interaction( method generate_response (line 53) | async def generate_response( method calculate_score (line 78) | async def calculate_score(self, instance_id: str, **kwargs) -> float: method finalize_interaction (line 87) | async def finalize_interaction(self, instance_id: str, **kwargs) -> None: FILE: siirl/execution/rollout_flow/multiturn/interactions/utils/interaction_registry.py function get_interaction_class (line 25) | def get_interaction_class(cls_name): function initialize_interactions_from_config (line 40) | def initialize_interactions_from_config(interaction_config_file): FILE: siirl/execution/rollout_flow/multiturn/tools/base_tool.py class BaseTool (line 22) | class BaseTool: method __init__ (line 34) | def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema): method get_openai_tool_schema (line 41) | def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema: method create (line 44) | async def create(self, instance_id: Optional[str] = None, **kwargs) ->... method execute (line 58) | async def execute(self, instance_id: str, parameters: dict[str, Any], ... method calc_reward (line 72) | async def calc_reward(self, instance_id: str, **kwargs) -> float: method release (line 83) | async def release(self, instance_id: str, **kwargs) -> None: FILE: siirl/execution/rollout_flow/multiturn/tools/geo3k_tool.py class Geo3kTool (line 28) | class Geo3kTool(BaseTool): method __init__ (line 37) | def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema): method get_openai_tool_schema (line 60) | def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema: method create (line 63) | async def create(self, instance_id: Optional[str] = None, ground_truth... method execute (line 73) | async def execute(self, instance_id: str, parameters: dict[str, Any], ... method calc_reward (line 87) | async def calc_reward(self, instance_id: str, **kwargs) -> float: method release (line 95) | async def release(self, instance_id: str, **kwargs) -> None: FILE: siirl/execution/rollout_flow/multiturn/tools/gsm8k_tool.py class Gsm8kTool (line 29) | class Gsm8kTool(BaseTool): method __init__ (line 39) | def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema): method get_openai_tool_schema (line 62) | def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema: method create (line 65) | async def create(self, instance_id: Optional[str] = None, ground_truth... method execute (line 75) | async def execute(self, instance_id: str, parameters: dict[str, Any], ... method calc_reward (line 93) | async def calc_reward(self, instance_id: str, **kwargs) -> float: method release (line 102) | async def release(self, instance_id: str, **kwargs) -> None: FILE: siirl/execution/rollout_flow/multiturn/tools/mcp_base_tool.py class MCPBaseTool (line 29) | class MCPBaseTool(BaseTool): method __init__ (line 30) | def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema): method get_openai_tool_schema (line 38) | def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema: method create (line 42) | async def create(self, instance_id: Optional[str] = None, **kwargs) ->... method _call_tool (line 59) | async def _call_tool(self, instance_id, parameters) -> Tuple[str, dict]: method execute (line 75) | async def execute(self, instance_id: str, parameters: dict[str, Any], ... method calc_reward (line 102) | async def calc_reward(self, instance_id: str, **kwargs) -> str: method release (line 105) | async def release(self, instance_id: str, **kwargs) -> None: method _parse_tool_result (line 109) | def _parse_tool_result(self, content: list) -> Tuple[str, dict]: FILE: siirl/execution/rollout_flow/multiturn/tools/mcp_search_tool.py class MCPSearchTool (line 26) | class MCPSearchTool(MCPBaseTool): method __init__ (line 27) | def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema): method _parse_tool_result (line 30) | def _parse_tool_result(self, content: list) -> Tuple[str, dict]: FILE: siirl/execution/rollout_flow/multiturn/tools/sandbox_fusion_tools.py class PoolMode (line 37) | class PoolMode(Enum): class TokenBucketWorker (line 43) | class TokenBucketWorker: method __init__ (line 44) | def __init__(self, rate_limit: int): method acquire (line 51) | def acquire(self): method release (line 56) | def release(self): method get_current_count (line 60) | def get_current_count(self): class ExecutionWorker (line 64) | class ExecutionWorker: method __init__ (line 65) | def __init__(self, enable_global_rate_limit=True, rate_limit=10): method _init_rate_limit (line 68) | def _init_rate_limit(self, rate_limit): method ping (line 73) | def ping(self): method execute (line 76) | def execute(self, fn: Callable[..., T], *fn_args, **fn_kwargs) -> T: function init_execution_pool (line 87) | def init_execution_pool( class SandboxFusionTool (line 101) | class SandboxFusionTool(BaseTool): method __init__ (line 111) | def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema): method get_openai_tool_schema (line 152) | def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema: method create (line 155) | async def create(self, instance_id: Optional[str] = None, ground_truth... method execute (line 165) | async def execute(self, instance_id: str, parameters: dict[str, Any], ... method execute_code (line 176) | def execute_code(self, instance_id, code, timeout=30, language="python"): method calc_reward (line 188) | async def calc_reward(self, instance_id: str, **kwargs) -> str: method release (line 191) | async def release(self, instance_id: str, **kwargs) -> None: FILE: siirl/execution/rollout_flow/multiturn/tools/schemas.py class OpenAIFunctionPropertySchema (line 21) | class OpenAIFunctionPropertySchema(BaseModel): class OpenAIFunctionParametersSchema (line 29) | class OpenAIFunctionParametersSchema(BaseModel): class OpenAIFunctionSchema (line 37) | class OpenAIFunctionSchema(BaseModel): class OpenAIFunctionToolSchema (line 46) | class OpenAIFunctionToolSchema(BaseModel): class OpenAIFunctionParsedSchema (line 53) | class OpenAIFunctionParsedSchema(BaseModel): class OpenAIFunctionCallSchema (line 60) | class OpenAIFunctionCallSchema(BaseModel): method from_openai_function_parsed_schema (line 67) | def from_openai_function_parsed_schema( class OpenAIFunctionToolCall (line 84) | class OpenAIFunctionToolCall(BaseModel): FILE: siirl/execution/rollout_flow/multiturn/tools/search_tool.py class PoolMode (line 39) | class PoolMode(Enum): class TokenBucketWorker (line 47) | class TokenBucketWorker: method __init__ (line 50) | def __init__(self, rate_limit: int): method acquire (line 56) | def acquire(self): method release (line 62) | def release(self): method get_current_count (line 67) | def get_current_count(self): class SearchExecutionWorker (line 72) | class SearchExecutionWorker: method __init__ (line 75) | def __init__(self, enable_global_rate_limit=True, rate_limit=10): method _init_rate_limit (line 78) | def _init_rate_limit(self, rate_limit): method ping (line 82) | def ping(self): method execute (line 86) | def execute(self, fn: Callable[..., T], *fn_args, **fn_kwargs) -> T: function init_search_execution_pool (line 101) | def init_search_execution_pool( class SearchTool (line 115) | class SearchTool(BaseTool): method __init__ (line 130) | def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema): method get_openai_tool_schema (line 182) | def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema: method create (line 186) | async def create(self, instance_id: Optional[str] = None, **kwargs) ->... method execute_search (line 203) | def execute_search(self, instance_id: str, query_list: list, retrieval... method execute (line 226) | async def execute(self, instance_id: str, parameters: dict[str, Any], ... method calc_reward (line 270) | async def calc_reward(self, instance_id: str, **kwargs) -> str: method release (line 273) | async def release(self, instance_id: str, **kwargs) -> None: FILE: siirl/execution/rollout_flow/multiturn/tools/utils/mcp_clients/McpClientManager.py class MCPClientManager (line 28) | class MCPClientManager: method initialize (line 35) | async def initialize(self, config_path, rate_limit: float = 10.0): method call_tool (line 58) | async def call_tool(self, tool_name, parameters, timeout): method fetch_tool_schemas (line 67) | async def fetch_tool_schemas(self, tool_selected_list: list[str]) -> l... method get_client_with_tool_name (line 82) | def get_client_with_tool_name(self, tool_name: str): method _load_config (line 85) | def _load_config(self, file: str) -> dict[str, Any]: FILE: siirl/execution/rollout_flow/multiturn/tools/utils/mcp_clients/utils.py class TokenBucket (line 24) | class TokenBucket: method __init__ (line 25) | def __init__(self, rate_limit: float): method acquire (line 31) | def acquire(self) -> bool: function mcp2openai (line 45) | def mcp2openai(mcp_tool: Tool) -> dict: FILE: siirl/execution/rollout_flow/multiturn/tools/utils/search_r1_like_utils.py function call_search_api (line 34) | def call_search_api( function _passages2string (line 130) | def _passages2string(retrieval_result): function perform_single_search_batch (line 141) | def perform_single_search_batch( FILE: siirl/execution/rollout_flow/multiturn/tools/utils/tool_registry.py class ToolType (line 29) | class ToolType(Enum): function initialize_mcp_tool (line 34) | async def initialize_mcp_tool(tool_cls, tool_config) -> list: function get_tool_class (line 66) | def get_tool_class(cls_name): function initialize_tools_from_config (line 80) | def initialize_tools_from_config(tools_config_file): FILE: siirl/execution/scheduler/enums.py class AdvantageEstimator (line 18) | class AdvantageEstimator(str, Enum): class WorkflowType (line 29) | class WorkflowType(str, Enum): class Role (line 35) | class Role(Enum): class AlgorithmType (line 49) | class AlgorithmType(Enum): FILE: siirl/execution/scheduler/graph_updater.py function unflatten_dict_with_omegaconf (line 35) | def unflatten_dict_with_omegaconf(flat_dict: Dict[str, Any]) -> Dict[str... function update_task_graph_node_configs (line 57) | def update_task_graph_node_configs(workerflow_taskgraph: TaskGraph, basi... function display_node_config (line 131) | def display_node_config(workerflow_taskgraph: TaskGraph) -> None: FILE: siirl/execution/scheduler/launch.py class RayTrainer (line 31) | class RayTrainer: method __init__ (line 42) | def __init__(self, config: SiiRLArguments, process_group_manager: Proc... method _check_mutually_exclusive (line 94) | def _check_mutually_exclusive(self, component_config: Dict[str, Any], ... method validate_actor_config (line 118) | def validate_actor_config(self, node: Node, actor_conf: ActorArguments... method validate_reference_config (line 154) | def validate_reference_config(self, node: Node, reference_conf: RefArg... method validate_rollout_config (line 169) | def validate_rollout_config(self, node: Node, rollout_conf: RolloutArg... method validate_critic_config (line 194) | def validate_critic_config(self, node: Node, critic_conf: CriticArgume... method validate_reward_model_config (line 219) | def validate_reward_model_config(self, node: Node, reward_model_conf: ... method validate_configurations_for_task_graph (line 227) | def validate_configurations_for_task_graph(self, task_graph: TaskGraph... method _validate_config (line 305) | def _validate_config(self): method init_workers (line 310) | def init_workers(self): method start_workers (line 328) | def start_workers(self): FILE: siirl/execution/scheduler/process_group_manager.py class ProcessGroupManager (line 22) | class ProcessGroupManager: method __init__ (line 43) | def __init__( method _clear_internal_mappings (line 91) | def _clear_internal_mappings(self): method _collect_initial_topology_info (line 99) | def _collect_initial_topology_info( method _aggregate_ranks_for_nodes (line 145) | def _aggregate_ranks_for_nodes( method _populate_node_rank_mappings (line 165) | def _populate_node_rank_mappings(self, node_id_to_final_ranks: Dict[st... method _define_process_groups (line 170) | def _define_process_groups(self) -> Dict[Tuple[int, ...], str]: method _populate_final_node_and_type_assignments (line 188) | def _populate_final_node_and_type_assignments( method _populate_subgraph_node_type_process_group_mapping (line 205) | def _populate_subgraph_node_type_process_group_mapping( method _compute_group_configurations (line 219) | def _compute_group_configurations(self): method get_group_spec (line 252) | def get_group_spec(self, group_name: str) -> Optional[Dict[str, Any]]: method get_all_specs (line 257) | def get_all_specs(self) -> Dict[str, Dict[str, Any]]: method get_node_assignment (line 261) | def get_node_assignment(self, node_id: str) -> Optional[Dict[str, Any]]: method get_process_groups_for_node_type (line 270) | def get_process_groups_for_node_type(self, node_type_value: str) -> Se... method get_process_group_for_node_type_in_subgraph (line 277) | def get_process_group_for_node_type_in_subgraph(self, graph_id: str, n... function _format_ranks_for_logging (line 290) | def _format_ranks_for_logging(ranks: Optional[List[int]], detailed_print... function _log_group_specs_report (line 313) | def _log_group_specs_report(pgm: ProcessGroupManager, detailed_printing:... function _log_node_assignments_report (line 328) | def _log_node_assignments_report( function _log_global_type_mappings_report (line 364) | def _log_global_type_mappings_report(pgm: ProcessGroupManager, node_type... function _log_subgraph_mappings_report (line 393) | def _log_subgraph_mappings_report( function log_process_group_manager_details (line 434) | def log_process_group_manager_details( FILE: siirl/execution/scheduler/ray_actor_manager.py class DistributedEnv (line 41) | class DistributedEnv(Enum): class RayActorManager (line 62) | class RayActorManager(WorkerGroup): method __init__ (line 79) | def __init__( method _initialize_workers (line 140) | def _initialize_workers(self, resource_pool: RayResourcePool, bin_pack... method _get_register_center_and_master_info (line 182) | def _get_register_center_and_master_info(self) -> Tuple[str, str]: method _create_worker_actor (line 225) | def _create_worker_actor( method _is_worker_alive (line 288) | def _is_worker_alive(self, worker: ActorHandle) -> bool: method _invoke_on_worker (line 306) | def _invoke_on_worker(self, worker: ActorHandle, method_name: str, *ar... method map_sync (line 311) | def map_sync(self, method_name: str, *args: Any, **kwargs: Any) -> Lis... method map_async (line 327) | def map_async(self, method_name: str, *args: Any, **kwargs: Any) -> Li... method map (line 362) | def map(self, method_name: str, *args: Any, **kwargs: Any) -> List[ray... method worker_names (line 369) | def worker_names(self) -> List[str]: method master_address (line 374) | def master_address(self) -> Optional[str]: method master_port (line 379) | def master_port(self) -> Optional[str]: method workers (line 384) | def workers(self) -> List[ActorHandle]: method world_size (line 389) | def world_size(self) -> int: FILE: siirl/execution/scheduler/resource_manager.py class ResourcePoolManager (line 32) | class ResourcePoolManager: method create_resource_pool (line 40) | def create_resource_pool(self): method get_resource_pool (line 51) | def get_resource_pool(self, resource_pool_name: str) -> RayResourcePool: method get_n_gpus (line 54) | def get_n_gpus(self) -> int: method _check_resource_available (line 58) | def _check_resource_available(self, timeout=60, interval=1): method _verify_placement_possible (line 105) | def _verify_placement_possible(self, available_gpus_per_node: dict): FILE: siirl/execution/scheduler/reward.py function load_custom_reward_function (line 44) | def load_custom_reward_function(config: SiiRLArguments) -> Optional[Call... function create_reward_manager (line 101) | def create_reward_manager( function compute_reward (line 192) | def compute_reward(data: TensorDict, reward_fn: Callable) -> Tuple[Rewar... FILE: siirl/execution/scheduler/task_scheduler.py function _parse_model_params_string (line 24) | def _parse_model_params_string(params_value: any) -> float: function estimate_graph_model_params (line 58) | def estimate_graph_model_params(task_graph: TaskGraph) -> float: class TaskScheduler (line 101) | class TaskScheduler: method __init__ (line 108) | def __init__(self, num_physical_nodes: int, gpus_per_node: int): method _reset_scheduler_state (line 134) | def _reset_scheduler_state(self): method _get_original_graph_id (line 147) | def _get_original_graph_id(self, task_graph: TaskGraph) -> str: method _apportion_workers_to_tasks (line 167) | def _apportion_workers_to_tasks(self, task_graphs_with_estimated_sizes... method schedule_and_assign_tasks (line 253) | def schedule_and_assign_tasks( method get_unique_assigned_task_graphs (line 424) | def get_unique_assigned_task_graphs(self) -> Dict[str, TaskGraph]: function _format_ranks_for_logging (line 447) | def _format_ranks_for_logging(ranks: Optional[List[int]], detailed_rank_... function log_schedule_assignments (line 468) | def log_schedule_assignments( function create_dummy_graph (line 572) | def create_dummy_graph(graph_id: str, num_nodes: int, model_params: any ... FILE: siirl/main_dag.py function _maybe_prepare_embodied_manifest (line 45) | def _maybe_prepare_embodied_manifest(siirl_args: SiiRLArguments) -> None: function load_pipeline (line 85) | def load_pipeline(siirl_args: SiiRLArguments) -> TaskGraph: class MainRunner (line 186) | class MainRunner: method run (line 196) | def run(self, siirl_args: SiiRLArguments) -> None: function main (line 266) | def main() -> None: FILE: siirl/models/embodied/openvla/configuration_prismatic.py class PrismaticConfig (line 72) | class PrismaticConfig(PretrainedConfig): method __init__ (line 76) | def __init__( class OpenVLAConfig (line 129) | class OpenVLAConfig(PrismaticConfig): method __init__ (line 132) | def __init__( FILE: siirl/models/embodied/openvla/modeling_prismatic.py function unpack_tuple (line 41) | def unpack_tuple(fn: Callable[[Any], Tuple[Any]]) -> Callable[[Any], Any]: function _ls_new_forward (line 52) | def _ls_new_forward(self, x: torch.Tensor) -> torch.Tensor: function ls_apply_patch (line 56) | def ls_apply_patch(ls_module: LayerScale): class PrismaticVisionBackbone (line 63) | class PrismaticVisionBackbone(nn.Module): method __init__ (line 64) | def __init__( method forward (line 114) | def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: class PrismaticProjector (line 127) | class PrismaticProjector(nn.Module): method __init__ (line 128) | def __init__(self, use_fused_vision_backbone: bool, vision_dim: int, l... method forward (line 146) | def forward(self, img_patches: torch.Tensor) -> torch.Tensor: class PrismaticCausalLMOutputWithPast (line 163) | class PrismaticCausalLMOutputWithPast(ModelOutput): class PrismaticPreTrainedModel (line 176) | class PrismaticPreTrainedModel(PreTrainedModel): method _init_weights (line 185) | def _init_weights(self, module: nn.Module) -> None: method _supports_sdpa (line 208) | def _supports_sdpa(self) -> bool: class PrismaticForConditionalGeneration (line 213) | class PrismaticForConditionalGeneration(PrismaticPreTrainedModel): method __init__ (line 214) | def __init__(self, config: PrismaticConfig) -> None: method get_input_embeddings (line 258) | def get_input_embeddings(self) -> nn.Module: method set_input_embeddings (line 261) | def set_input_embeddings(self, value: nn.Module) -> None: method get_output_embeddings (line 264) | def get_output_embeddings(self) -> nn.Module: method set_output_embeddings (line 267) | def set_output_embeddings(self, new_embeddings: nn.Module) -> None: method get_decoder (line 270) | def get_decoder(self) -> nn.Module: method set_decoder (line 273) | def set_decoder(self, decoder: nn.Module) -> None: method tie_weights (line 276) | def tie_weights(self) -> None: method resize_token_embeddings (line 279) | def resize_token_embeddings( method forward (line 291) | def forward( method prepare_inputs_for_generation (line 450) | def prepare_inputs_for_generation( method _reorder_cache (line 488) | def _reorder_cache(self, *args, **kwargs) -> Any: class OpenVLAForActionPrediction (line 492) | class OpenVLAForActionPrediction(PrismaticForConditionalGeneration): method __init__ (line 495) | def __init__(self, config: OpenVLAConfig) -> None: method predict_action (line 506) | def predict_action( method _check_unnorm_key (line 539) | def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_ke... method get_action_dim (line 554) | def get_action_dim(self, unnorm_key: Optional[str] = None) -> int: method get_action_stats (line 559) | def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict[s... FILE: siirl/models/embodied/openvla/processing_prismatic.py function letterbox_pad_transform (line 23) | def letterbox_pad_transform(image: Image.Image, padding_fill_value: Tupl... class PrismaticImageProcessor (line 32) | class PrismaticImageProcessor(ImageProcessingMixin): method __init__ (line 35) | def __init__( method apply_transform (line 129) | def apply_transform(self, img: Image.Image) -> torch.Tensor: method preprocess (line 148) | def preprocess( method __call__ (line 172) | def __call__(self, images: Union[Image.Image, List[Image.Image]], **kw... class PrismaticProcessor (line 178) | class PrismaticProcessor(ProcessorMixin): method __init__ (line 183) | def __init__( method __call__ (line 190) | def __call__( method batch_decode (line 224) | def batch_decode( method decode (line 238) | def decode( method model_input_names (line 253) | def model_input_names(self) -> List[str]: FILE: siirl/models/embodied/openvla_oft/configuration_prismatic.py class PrismaticConfig (line 72) | class PrismaticConfig(PretrainedConfig): method __init__ (line 76) | def __init__( class OpenVLAConfig (line 129) | class OpenVLAConfig(PrismaticConfig): method __init__ (line 132) | def __init__( FILE: siirl/models/embodied/openvla_oft/constants.py class NormalizationType (line 17) | class NormalizationType(str, Enum): function detect_robot_platform (line 49) | def detect_robot_platform(): FILE: siirl/models/embodied/openvla_oft/modeling_prismatic.py function unpack_tuple (line 45) | def unpack_tuple(fn: Callable[[Any], Tuple[Any]]) -> Callable[[Any], Any]: function _ls_new_forward (line 56) | def _ls_new_forward(self, x: torch.Tensor) -> torch.Tensor: function ls_apply_patch (line 60) | def ls_apply_patch(ls_module: LayerScale): class PrismaticVisionBackbone (line 67) | class PrismaticVisionBackbone(nn.Module): method __init__ (line 75) | def __init__( method _create_featurizer (line 115) | def _create_featurizer(self, model_id: str, img_size: int, act_layer: ... method _patch_layer_scales (line 141) | def _patch_layer_scales(self) -> None: method get_num_patches (line 159) | def get_num_patches(self) -> int: method get_num_images_in_input (line 168) | def get_num_images_in_input(self) -> int: method set_num_images_in_input (line 177) | def set_num_images_in_input(self, num_images_in_input: int) -> None: method forward (line 186) | def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: class PrismaticProjector (line 231) | class PrismaticProjector(nn.Module): method __init__ (line 232) | def __init__(self, use_fused_vision_backbone: bool, vision_dim: int, l... method forward (line 250) | def forward(self, img_patches: torch.Tensor) -> torch.Tensor: class PrismaticCausalLMOutputWithPast (line 267) | class PrismaticCausalLMOutputWithPast(ModelOutput): class PrismaticPreTrainedModel (line 280) | class PrismaticPreTrainedModel(PreTrainedModel): method _init_weights (line 289) | def _init_weights(self, module: nn.Module) -> None: method _supports_sdpa (line 312) | def _supports_sdpa(self) -> bool: class PrismaticForConditionalGeneration (line 317) | class PrismaticForConditionalGeneration(PrismaticPreTrainedModel): method __init__ (line 318) | def __init__(self, config: PrismaticConfig) -> None: method get_input_embeddings (line 363) | def get_input_embeddings(self) -> nn.Module: method set_input_embeddings (line 366) | def set_input_embeddings(self, value: nn.Module) -> None: method get_output_embeddings (line 369) | def get_output_embeddings(self) -> nn.Module: method set_output_embeddings (line 372) | def set_output_embeddings(self, new_embeddings: nn.Module) -> None: method get_decoder (line 375) | def get_decoder(self) -> nn.Module: method set_decoder (line 378) | def set_decoder(self, decoder: nn.Module) -> None: method tie_weights (line 381) | def tie_weights(self) -> None: method resize_token_embeddings (line 384) | def resize_token_embeddings( method _replace_input_embeddings (line 395) | def _replace_input_embeddings(self, input_embeddings, all_actions_mask... method _process_action_masks (line 431) | def _process_action_masks(self, labels): method _process_vision_features (line 438) | def _process_vision_features(self, pixel_values, language_embeddings=N... method _process_proprio_features (line 449) | def _process_proprio_features(self, projected_patch_embeddings, propri... method _build_multimodal_attention (line 461) | def _build_multimodal_attention(self, input_embeddings, projected_patc... method _build_multimodal_labels (line 486) | def _build_multimodal_labels(self, labels, projected_patch_embeddings): method prepare_inputs_for_generation (line 682) | def prepare_inputs_for_generation( method _reorder_cache (line 720) | def _reorder_cache(self, *args, **kwargs) -> Any: method _prepare_input_for_action_prediction_verl (line 723) | def _prepare_input_for_action_prediction_verl(self, input_ids, attenti... method _prepare_labels_for_action_prediction_verl (line 746) | def _prepare_labels_for_action_prediction_verl(self, labels, input_ids): method _verl_discrete_compute_logits (line 761) | def _verl_discrete_compute_logits( method forward (line 1077) | def forward( class OpenVLAForActionPrediction (line 1326) | class OpenVLAForActionPrediction(PrismaticForConditionalGeneration): method __init__ (line 1329) | def __init__(self, config: OpenVLAConfig) -> None: method _prepare_input_for_action_prediction (line 1340) | def _prepare_input_for_action_prediction(self, input_ids, attention_ma... method _prepare_labels_for_action_prediction (line 1363) | def _prepare_labels_for_action_prediction(self, labels, input_ids): method _unnormalize_actions (line 1378) | def _unnormalize_actions(self, normalized_actions, unnorm_key=None): method _run_diffusion_prediction (line 1399) | def _run_diffusion_prediction( method _regression_or_discrete_prediction (line 1485) | def _regression_or_discrete_prediction( method _verl_discrete_prediction (line 1552) | def _verl_discrete_prediction( method predict_action (line 1711) | def predict_action( method generate_action_verl (line 1827) | def generate_action_verl( method _check_unnorm_key (line 1972) | def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_ke... method get_action_dim (line 1988) | def get_action_dim(self, unnorm_key: Optional[str] = None) -> int: method get_action_stats (line 1993) | def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict[s... FILE: siirl/models/embodied/openvla_oft/processing_prismatic.py function letterbox_pad_transform (line 23) | def letterbox_pad_transform(image: Image.Image, padding_fill_value: Tupl... class PrismaticImageProcessor (line 32) | class PrismaticImageProcessor(ImageProcessingMixin): method __init__ (line 35) | def __init__( method apply_transform (line 128) | def apply_transform(self, img: Image.Image) -> torch.Tensor: method preprocess (line 147) | def preprocess( method __call__ (line 169) | def __call__(self, images: Union[Image.Image, List[Image.Image]], **kw... class PrismaticProcessor (line 175) | class PrismaticProcessor(ProcessorMixin): method __init__ (line 180) | def __init__( method __call__ (line 187) | def __call__( method batch_decode (line 219) | def batch_decode( method decode (line 233) | def decode( method model_input_names (line 248) | def model_input_names(self) -> List[str]: FILE: siirl/models/embodied/openvla_oft/train_utils.py function get_current_action_mask (line 8) | def get_current_action_mask(token_ids): function get_next_actions_mask (line 25) | def get_next_actions_mask(token_ids): function compute_token_accuracy (line 42) | def compute_token_accuracy(predicted_token_ids, ground_truth_token_ids, ... function compute_actions_l1_loss (line 48) | def compute_actions_l1_loss(action_tokenizer, predicted_token_ids, groun... FILE: siirl/models/llama/megatron/checkpoint_utils/llama_loader.py function _megatron_calc_layer_map (line 23) | def _megatron_calc_layer_map(config): function load_state_dict_to_megatron_llama (line 53) | def load_state_dict_to_megatron_llama(state_dict, wrapped_models, config... FILE: siirl/models/llama/megatron/checkpoint_utils/llama_loader_depracated.py function _megatron_calc_layer_map (line 23) | def _megatron_calc_layer_map(config): function load_state_dict_to_megatron_llama (line 53) | def load_state_dict_to_megatron_llama(state_dict, wrapped_models, config... FILE: siirl/models/llama/megatron/checkpoint_utils/llama_saver.py function _megatron_calc_global_rank (line 28) | def _megatron_calc_global_rank(tp_rank: int = 0, dp_rank: int = 0, pp_ra... function _megatron_calc_layer_map (line 39) | def _megatron_calc_layer_map(config): function merge_megatron_ckpt_llama (line 67) | def merge_megatron_ckpt_llama(wrapped_models, config, dtype, is_value_mo... FILE: siirl/models/llama/megatron/layers/parallel_attention.py class LlamaRotaryEmbedding (line 38) | class LlamaRotaryEmbedding(nn.Module): method __init__ (line 39) | def __init__(self, dim, max_position_embeddings=2048, base=10000, devi... method _set_cos_sin_cache (line 51) | def _set_cos_sin_cache(self, seq_len, device, dtype): method forward (line 61) | def forward(self, x, seq_len=None): class LlamaLinearScalingRotaryEmbedding (line 72) | class LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding): method __init__ (line 75) | def __init__(self, dim, max_position_embeddings=2048, base=10000, devi... method _set_cos_sin_cache (line 79) | def _set_cos_sin_cache(self, seq_len, device, dtype): class LlamaDynamicNTKScalingRotaryEmbedding (line 91) | class LlamaDynamicNTKScalingRotaryEmbedding(LlamaRotaryEmbedding): method __init__ (line 94) | def __init__(self, dim, max_position_embeddings=2048, base=10000, devi... method _set_cos_sin_cache (line 98) | def _set_cos_sin_cache(self, seq_len, device, dtype): class LlamaLlama3ScalingRotaryEmbedding (line 115) | class LlamaLlama3ScalingRotaryEmbedding(LlamaRotaryEmbedding): method __init__ (line 116) | def __init__(self, dim, config, max_position_embeddings=2048, base=100... function rotate_half (line 142) | def rotate_half(x): function apply_rotary_pos_emb (line 149) | def apply_rotary_pos_emb(q, k, cos, sin, position_ids): function repeat_kv (line 157) | def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: class ParallelLlamaAttention (line 169) | class ParallelLlamaAttention(nn.Module): method __init__ (line 172) | def __init__(self, config: LlamaConfig, megatron_config: ModelParallel... method _init_rope (line 232) | def _init_rope(self): method _shape (line 267) | def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): method forward (line 270) | def forward( function apply_rotary_pos_emb_rmpad (line 326) | def apply_rotary_pos_emb_rmpad(q, k, cos, sin, position_ids, indices, se... function apply_rotary_pos_emb_rmpad_flash (line 344) | def apply_rotary_pos_emb_rmpad_flash(q, k, cos, sin, cu_seqlens, max_seq... class ParallelLlamaAttentionRmPad (line 350) | class ParallelLlamaAttentionRmPad(ParallelLlamaAttention): method forward (line 351) | def forward( FILE: siirl/models/llama/megatron/layers/parallel_decoder.py class ParallelLlamaDecoderLayer (line 35) | class ParallelLlamaDecoderLayer(nn.Module): method __init__ (line 36) | def __init__(self, config: LlamaConfig, megatron_config: ModelParallel... method forward (line 47) | def forward( class ParallelLlamaDecoderLayerRmPad (line 102) | class ParallelLlamaDecoderLayerRmPad(nn.Module): method __init__ (line 103) | def __init__(self, config: LlamaConfig, megatron_config: ModelParallel... method forward (line 114) | def forward( FILE: siirl/models/llama/megatron/layers/parallel_linear.py class QKVParallelLinear (line 20) | class QKVParallelLinear(tensor_parallel.ColumnParallelLinear): method __init__ (line 21) | def __init__( class MergedColumnParallelLinear (line 54) | class MergedColumnParallelLinear(tensor_parallel.ColumnParallelLinear): method __init__ (line 55) | def __init__( class LinearForLastLayer (line 82) | class LinearForLastLayer(torch.nn.Linear): method __init__ (line 83) | def __init__( method forward (line 96) | def forward( FILE: siirl/models/llama/megatron/layers/parallel_mlp.py class ParallelLlamaMLP (line 30) | class ParallelLlamaMLP(nn.Module): method __init__ (line 31) | def __init__(self, config, megatron_config: ModelParallelConfig = None... method forward (line 71) | def forward(self, x): FILE: siirl/models/llama/megatron/layers/parallel_rmsnorm.py class ParallelLlamaRMSNorm (line 26) | class ParallelLlamaRMSNorm(nn.Module): method __init__ (line 27) | def __init__(self, config: LlamaConfig, megatron_config: ModelParallel... method forward (line 41) | def forward(self, hidden_states): FILE: siirl/models/llama/megatron/modeling_llama_megatron.py function _make_causal_mask (line 47) | def _make_causal_mask(input_ids_shape: torch.Size, dtype: torch.dtype, d... function _expand_mask (line 60) | def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Option... class ParallelLlamaModel (line 74) | class ParallelLlamaModel(nn.Module): method __init__ (line 82) | def __init__(self, config: LlamaConfig, megatron_config: ModelParallel... method _prepare_decoder_attention_mask (line 97) | def _prepare_decoder_attention_mask(self, attention_mask, input_shape,... method forward (line 115) | def forward( class ParallelLlamaForCausalLM (line 153) | class ParallelLlamaForCausalLM(nn.Module): method __init__ (line 154) | def __init__(self, config: LlamaConfig, megatron_config: ModelParallel... method forward (line 174) | def forward( class ParallelLlamaModelRmPad (line 215) | class ParallelLlamaModelRmPad(nn.Module): method __init__ (line 223) | def __init__(self, config: LlamaConfig, megatron_config: ModelParallel... method forward (line 238) | def forward( class ParallelLlamaForCausalLMRmPad (line 281) | class ParallelLlamaForCausalLMRmPad(nn.Module): method __init__ (line 282) | def __init__(self, config: LlamaConfig, megatron_config: ModelParallel... method _init_head (line 290) | def _init_head(self, config): method _forward_head (line 304) | def _forward_head(self, hidden_states): method forward (line 311) | def forward( class ParallelLlamaForValueRmPad (line 369) | class ParallelLlamaForValueRmPad(ParallelLlamaForCausalLMRmPad): method _init_head (line 370) | def _init_head(self, config): method _forward_head (line 379) | def _forward_head(self, hidden_states): method forward (line 386) | def forward( class ParallelLlamaModelRmPadPP (line 402) | class ParallelLlamaModelRmPadPP(nn.Module): method __init__ (line 412) | def __init__(self, config: LlamaConfig, megatron_config: ModelParallel... method set_input_tensor (line 454) | def set_input_tensor(self, input_tensor): method forward (line 464) | def forward( class ParallelLlamaForCausalLMRmPadPP (line 515) | class ParallelLlamaForCausalLMRmPadPP(nn.Module): method __init__ (line 516) | def __init__( method set_input_tensor (line 536) | def set_input_tensor(self, input_tensor): method _init_head (line 547) | def _init_head(self, config): method _forward_head (line 561) | def _forward_head(self, hidden_states): method forward (line 569) | def forward( class ParallelLlamaForValueRmPadPP (line 633) | class ParallelLlamaForValueRmPadPP(ParallelLlamaForCausalLMRmPadPP): method _init_head (line 634) | def _init_head(self, config): method _forward_head (line 643) | def _forward_head(self, hidden_states): method forward (line 650) | def forward( FILE: siirl/models/loader.py class TokenizerModule (line 21) | class TokenizerModule(TypedDict): function _get_init_kwargs (line 26) | def _get_init_kwargs(model_args: "ModelArguments") -> Dict[str, Any]: function set_pad_token_id (line 41) | def set_pad_token_id(tokenizer): function load_tokenizer (line 56) | def load_tokenizer( FILE: siirl/models/mcore/config_converter.py function _get_base_transformer_config (line 26) | def _get_base_transformer_config(hf_config: PretrainedConfig, dtype: tor... function hf_to_mcore_config_dense (line 88) | def hf_to_mcore_config_dense(hf_config: PretrainedConfig, dtype: torch.d... function hf_to_mcore_config_qwen2moe (line 97) | def hf_to_mcore_config_qwen2moe(hf_config: PretrainedConfig, dtype: torc... function hf_to_mcore_config_mixtral (line 128) | def hf_to_mcore_config_mixtral(hf_config: PretrainedConfig, dtype: torch... function hf_to_mcore_config_qwen3moe (line 158) | def hf_to_mcore_config_qwen3moe(hf_config: PretrainedConfig, dtype: torc... function hf_to_mcore_config_dpskv3 (line 187) | def hf_to_mcore_config_dpskv3(hf_config: PretrainedConfig, dtype: torch.... function hf_to_mcore_config_qwen2_5_vl (line 277) | def hf_to_mcore_config_qwen2_5_vl(hf_config: PretrainedConfig, dtype: to... function hf_to_mcore_config_llama4 (line 282) | def hf_to_mcore_config_llama4(hf_config: PretrainedConfig, dtype: torch.... FILE: siirl/models/mcore/loader.py function _megatron_calc_layer_map (line 26) | def _megatron_calc_layer_map(config): function load_state_dict_to_megatron_gptmodel (line 54) | def load_state_dict_to_megatron_gptmodel(state_dict, wrapped_models, con... FILE: siirl/models/mcore/model_forward.py function gptmodel_forward (line 22) | def gptmodel_forward( FILE: siirl/models/mcore/model_forward_fused.py function patch_fused_forward (line 36) | def patch_fused_forward(model: torch.nn.Module): function unpatch_fused_forward (line 47) | def unpatch_fused_forward(model: torch.nn.Module): function fused_forward_gptmodel (line 57) | def fused_forward_gptmodel( function _fused_GPTModel_forward (line 103) | def _fused_GPTModel_forward( FILE: siirl/models/mcore/model_initializer.py class BaseModelInitializer (line 26) | class BaseModelInitializer(ABC): method __init__ (line 29) | def __init__(self, tfconfig: TransformerConfig, hf_config: PretrainedC... method get_transformer_layer_spec (line 34) | def get_transformer_layer_spec(self): method get_rope_scaling_args (line 39) | def get_rope_scaling_args(self) -> dict: method initialize (line 48) | def initialize( class DenseModel (line 93) | class DenseModel(BaseModelInitializer): method get_transformer_layer_spec (line 96) | def get_transformer_layer_spec(self): class Qwen2MoEModel (line 101) | class Qwen2MoEModel(BaseModelInitializer): method get_transformer_layer_spec (line 104) | def get_transformer_layer_spec(self): method initialize (line 114) | def initialize(self, **kwargs): class MixtralModel (line 124) | class MixtralModel(BaseModelInitializer): method get_transformer_layer_spec (line 127) | def get_transformer_layer_spec(self): method initialize (line 132) | def initialize(self, **kwargs): class Qwen3MoEModel (line 141) | class Qwen3MoEModel(BaseModelInitializer): method get_transformer_layer_spec (line 144) | def get_transformer_layer_spec(self): method initialize (line 149) | def initialize(self, **kwargs): class DeepseekV3Model (line 159) | class DeepseekV3Model(BaseModelInitializer): method get_transformer_layer_spec (line 162) | def get_transformer_layer_spec(self): method get_rope_scaling_args (line 166) | def get_rope_scaling_args(self) -> dict: method initialize (line 171) | def initialize( class Qwen25VLModel (line 192) | class Qwen25VLModel(BaseModelInitializer): method get_transformer_layer_spec (line 195) | def get_transformer_layer_spec(self): FILE: siirl/models/mcore/patch_v012.py function apply_patch (line 20) | def apply_patch(): FILE: siirl/models/mcore/registry.py class SupportedModel (line 60) | class SupportedModel(Enum): function get_supported_model (line 137) | def get_supported_model(model_type: str) -> SupportedModel: function hf_to_mcore_config (line 145) | def hf_to_mcore_config(hf_config: PretrainedConfig, dtype: torch.dtype, ... function init_mcore_model (line 151) | def init_mcore_model( function get_mcore_forward_fn (line 183) | def get_mcore_forward_fn(hf_config: PretrainedConfig) -> Callable: function get_mcore_forward_fused_fn (line 191) | def get_mcore_forward_fused_fn(hf_config: PretrainedConfig) -> Callable: function get_mcore_weight_converter (line 199) | def get_mcore_weight_converter(hf_config: PretrainedConfig, dtype: torch... FILE: siirl/models/mcore/saver.py function _megatron_calc_global_rank (line 29) | def _megatron_calc_global_rank(tp_rank: int = 0, dp_rank: int = 0, pp_ra... function _megatron_calc_layer_map (line 48) | def _megatron_calc_layer_map(config): function merge_megatron_ckpt_gptmodel (line 76) | def merge_megatron_ckpt_gptmodel(wrapped_models, config, dtype, is_value... function merge_megatron_ckpt_gptmodel_qwen_moe (line 467) | def merge_megatron_ckpt_gptmodel_qwen_moe(wrapped_models, config, dtype,... function merge_megatron_ckpt_gptmodel_dpskv3 (line 471) | def merge_megatron_ckpt_gptmodel_dpskv3(wrapped_models, config, dtype, i... function merge_megatron_ckpt_gptmodel_mixtral (line 475) | def merge_megatron_ckpt_gptmodel_mixtral(wrapped_models, config, dtype, ... FILE: siirl/models/mcore/util.py function preprocess_packed_seqs (line 23) | def preprocess_packed_seqs( function postprocess_packed_seqs (line 105) | def postprocess_packed_seqs( function remove_left_padding (line 165) | def remove_left_padding( function recover_left_padding (line 206) | def recover_left_padding( function postprocess_packed_seqs_for_dict_output (line 228) | def postprocess_packed_seqs_for_dict_output( FILE: siirl/models/mcore/weight_converter.py class McoreToHFWeightConverterBase (line 25) | class McoreToHFWeightConverterBase: method __init__ (line 26) | def __init__(self, hf_config: PretrainedConfig, mcore_config: Transfor... method convert_param (line 30) | def convert_param(self, name: str, params_one_group: list[torch.Tensor... class McoreToHFWeightConverterDense (line 34) | class McoreToHFWeightConverterDense(McoreToHFWeightConverterBase): method _convert_attention_param (line 35) | def _convert_attention_param(self, name: str, params: list[torch.Tenso... method _convert_mlp_param (line 65) | def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) ->... method convert_param (line 86) | def convert_param(self, name: str, params_one_group: list[torch.Tensor... class McoreToHFWeightConverterQwen2Moe (line 103) | class McoreToHFWeightConverterQwen2Moe(McoreToHFWeightConverterDense): method _convert_mlp_param (line 104) | def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) ->... class McoreToHFWeightConverterDpskv3 (line 150) | class McoreToHFWeightConverterDpskv3(McoreToHFWeightConverterBase): method _convert_attention_param (line 151) | def _convert_attention_param(self, name: str, params: list[torch.Tenso... method _convert_mlp_param (line 190) | def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) ->... method _convert_mtp_param (line 260) | def _convert_mtp_param(self, name: str, params: list[torch.Tensor]) ->... method convert_param (line 277) | def convert_param(self, name: str, params_one_group: list[torch.Tensor... class McoreToHFWeightConverterMixtral (line 295) | class McoreToHFWeightConverterMixtral(McoreToHFWeightConverterDense): method _convert_mlp_param (line 296) | def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) ->... class McoreToHFWeightConverterQwen3Moe (line 319) | class McoreToHFWeightConverterQwen3Moe(McoreToHFWeightConverterDense): method _convert_mlp_param (line 320) | def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) ->... FILE: siirl/models/model_utils/visual.py class CompositeModel (line 25) | class CompositeModel: method get_projector (line 32) | def get_projector(self, module: "torch.nn.Module") -> "torch.nn.Module": function _register_composite_model (line 42) | def _register_composite_model( class LlavaMultiModalProjectorForYiVL (line 58) | class LlavaMultiModalProjectorForYiVL(torch.nn.Module): method __init__ (line 59) | def __init__(self, config: "LlavaConfig") -> None: method forward (line 72) | def forward(self, image_features: "torch.Tensor") -> "torch.Tensor": class LlavaMultiModalProjectorForYiVLForVLLM (line 92) | class LlavaMultiModalProjectorForYiVLForVLLM(LlavaMultiModalProjectorFor... method __init__ (line 93) | def __init__(self, vision_hidden_size: int, text_hidden_size: int, pro... function autocast_projector_dtype (line 103) | def autocast_projector_dtype(model: "PreTrainedModel", model_args: "Mode... function configure_visual_model (line 122) | def configure_visual_model(config: "PretrainedConfig") -> None: function get_forbidden_modules (line 135) | def get_forbidden_modules(config: "PretrainedConfig", finetuning_args: "... function get_image_seqlen (line 160) | def get_image_seqlen(config: "PretrainedConfig") -> int: function get_patch_size (line 177) | def get_patch_size(config: "PretrainedConfig", processor: "ProcessorMixi... function get_vision_feature_select_strategy (line 185) | def get_vision_feature_select_strategy(config: "PretrainedConfig", proce... function patch_target_modules (line 197) | def patch_target_modules( FILE: siirl/models/patcher.py function patch_tokenizer (line 39) | def patch_tokenizer(tokenizer: "PreTrainedTokenizer", model_args: "Model... function patch_processor (line 66) | def patch_processor( FILE: siirl/models/qwen2/megatron/checkpoint_utils/qwen2_loader.py function _megatron_calc_layer_map (line 22) | def _megatron_calc_layer_map(config): function load_state_dict_to_megatron_qwen2 (line 50) | def load_state_dict_to_megatron_qwen2(state_dict, wrapped_models, config... FILE: siirl/models/qwen2/megatron/checkpoint_utils/qwen2_loader_depracated.py function _megatron_calc_layer_map (line 23) | def _megatron_calc_layer_map(config): function load_state_dict_to_megatron_qwen2 (line 51) | def load_state_dict_to_megatron_qwen2(state_dict, wrapped_models, config... FILE: siirl/models/qwen2/megatron/checkpoint_utils/qwen2_saver.py function _megatron_calc_global_rank (line 28) | def _megatron_calc_global_rank(tp_rank: int = 0, dp_rank: int = 0, pp_ra... function _megatron_calc_layer_map (line 39) | def _megatron_calc_layer_map(config): function merge_megatron_ckpt_qwen2 (line 67) | def merge_megatron_ckpt_qwen2(wrapped_models, config, dtype, is_value_mo... FILE: siirl/models/qwen2/megatron/layers/parallel_attention.py class Qwen2RotaryEmbedding (line 42) | class Qwen2RotaryEmbedding(nn.Module): method __init__ (line 43) | def __init__(self, dim, max_position_embeddings=2048, base=10000, devi... method _set_cos_sin_cache (line 55) | def _set_cos_sin_cache(self, seq_len, device, dtype): method forward (line 65) | def forward(self, x, seq_len=None): class Qwen2LinearScalingRotaryEmbedding (line 76) | class Qwen2LinearScalingRotaryEmbedding(Qwen2RotaryEmbedding): method __init__ (line 79) | def __init__(self, dim, max_position_embeddings=2048, base=10000, devi... method _set_cos_sin_cache (line 83) | def _set_cos_sin_cache(self, seq_len, device, dtype): class Qwen2DynamicNTKScalingRotaryEmbedding (line 95) | class Qwen2DynamicNTKScalingRotaryEmbedding(Qwen2RotaryEmbedding): method __init__ (line 98) | def __init__(self, dim, max_position_embeddings=2048, base=10000, devi... method _set_cos_sin_cache (line 102) | def _set_cos_sin_cache(self, seq_len, device, dtype): function rotate_half (line 119) | def rotate_half(x): function apply_rotary_pos_emb (line 126) | def apply_rotary_pos_emb(q, k, cos, sin, position_ids): function repeat_kv (line 134) | def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: class ParallelQwen2Attention (line 146) | class ParallelQwen2Attention(nn.Module): method __init__ (line 149) | def __init__(self, config: Qwen2Config, megatron_config: ModelParallel... method _init_rope (line 211) | def _init_rope(self): method _shape (line 218) | def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): method forward (line 221) | def forward( function apply_rotary_pos_emb_rmpad (line 272) | def apply_rotary_pos_emb_rmpad(q, k, cos, sin, position_ids, indices, se... function apply_rotary_pos_emb_rmpad_flash (line 290) | def apply_rotary_pos_emb_rmpad_flash(q, k, cos, sin, cu_seqlens, max_seq... class ParallelQwen2AttentionRmPad (line 296) | class ParallelQwen2AttentionRmPad(ParallelQwen2Attention): method forward (line 297) | def forward( FILE: siirl/models/qwen2/megatron/layers/parallel_decoder.py class ParallelQwen2DecoderLayer (line 35) | class ParallelQwen2DecoderLayer(nn.Module): method __init__ (line 36) | def __init__(self, config: Qwen2Config, megatron_config: ModelParallel... method forward (line 47) | def forward( class ParallelQwen2DecoderLayerRmPad (line 102) | class ParallelQwen2DecoderLayerRmPad(nn.Module): method __init__ (line 103) | def __init__(self, config: Qwen2Config, megatron_config: ModelParallel... method forward (line 114) | def forward( FILE: siirl/models/qwen2/megatron/layers/parallel_linear.py class QKVParallelLinear (line 20) | class QKVParallelLinear(tensor_parallel.ColumnParallelLinear): method __init__ (line 21) | def __init__( class MergedColumnParallelLinear (line 54) | class MergedColumnParallelLinear(tensor_parallel.ColumnParallelLinear): method __init__ (line 55) | def __init__( FILE: siirl/models/qwen2/megatron/layers/parallel_mlp.py class ParallelQwen2MLP (line 30) | class ParallelQwen2MLP(nn.Module): method __init__ (line 31) | def __init__(self, config, megatron_config: ModelParallelConfig = None... method forward (line 71) | def forward(self, x): FILE: siirl/models/qwen2/megatron/layers/parallel_rmsnorm.py class ParallelQwen2RMSNorm (line 26) | class ParallelQwen2RMSNorm(nn.Module): method __init__ (line 27) | def __init__(self, config: Qwen2Config, megatron_config: ModelParallel... method forward (line 41) | def forward(self, hidden_states): FILE: siirl/models/qwen2/megatron/modeling_qwen2_megatron.py function _make_causal_mask (line 48) | def _make_causal_mask(input_ids_shape: torch.Size, dtype: torch.dtype, d... function _expand_mask (line 61) | def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Option... class ParallelQwen2Model (line 75) | class ParallelQwen2Model(nn.Module): method __init__ (line 83) | def __init__(self, config: Qwen2Config, megatron_config: ModelParallel... method _prepare_decoder_attention_mask (line 98) | def _prepare_decoder_attention_mask(self, attention_mask, input_shape,... method forward (line 116) | def forward( class ParallelQwen2ForCausalLM (line 154) | class ParallelQwen2ForCausalLM(nn.Module): method __init__ (line 155) | def __init__(self, config: Qwen2Config, megatron_config: ModelParallel... method forward (line 175) | def forward( class ParallelQwen2ModelRmPad (line 216) | class ParallelQwen2ModelRmPad(nn.Module): method __init__ (line 224) | def __init__(self, config: Qwen2Config, megatron_config: ModelParallel... method forward (line 239) | def forward( class ParallelQwen2ForCausalLMRmPad (line 282) | class ParallelQwen2ForCausalLMRmPad(nn.Module): method __init__ (line 283) | def __init__(self, config: Qwen2Config, megatron_config: ModelParallel... method _init_head (line 291) | def _init_head(self, config: Qwen2Config): method _forward_head (line 305) | def _forward_head(self, hidden_states): method forward (line 312) | def forward( class ParallelQwen2ForValueRmPad (line 370) | class ParallelQwen2ForValueRmPad(ParallelQwen2ForCausalLMRmPad): method _init_head (line 371) | def _init_head(self, config): method _forward_head (line 380) | def _forward_head(self, hidden_states): method forward (line 387) | def forward( class ParallelQwen2ModelRmPadPP (line 403) | class ParallelQwen2ModelRmPadPP(nn.Module): method __init__ (line 413) | def __init__(self, config: Qwen2Config, megatron_config: ModelParallel... method set_input_tensor (line 454) | def set_input_tensor(self, input_tensor): method forward (line 464) | def forward( class ParallelQwen2ForCausalLMRmPadPP (line 515) | class ParallelQwen2ForCausalLMRmPadPP(nn.Module): method __init__ (line 516) | def __init__( method set_input_tensor (line 537) | def set_input_tensor(self, input_tensor): method _init_head (line 548) | def _init_head(self, config): method setup_embeddings_and_output_layer (line 563) | def setup_embeddings_and_output_layer(self) -> None: method shared_embedding_or_output_weight (line 602) | def shared_embedding_or_output_weight(self) -> torch.Tensor: method _forward_head (line 609) | def _forward_head(self, hidden_states): method forward (line 620) | def forward( class ParallelQwen2ForValueRmPadPP (line 683) | class ParallelQwen2ForValueRmPadPP(ParallelQwen2ForCausalLMRmPadPP): method _init_head (line 684) | def _init_head(self, config): method _forward_head (line 693) | def _forward_head(self, hidden_states): method forward (line 700) | def forward( FILE: siirl/models/registry.py class ModelRegistry (line 39) | class ModelRegistry: method load_model_cls (line 41) | def load_model_cls(model_arch: str, value=False) -> Optional[Type[nn.M... method get_supported_archs (line 57) | def get_supported_archs() -> List[str]: FILE: siirl/models/transformers/internvl.py class SeparatorStyle (line 55) | class SeparatorStyle(IntEnum): class Conversation (line 79) | class Conversation: method get_prompt (line 103) | def get_prompt(self) -> str: method set_system_message (line 289) | def set_system_message(self, system_message: str): method append_message (line 293) | def append_message(self, role: str, message: str): method update_last_message (line 297) | def update_last_message(self, message: str): method to_gradio_chatbot (line 305) | def to_gradio_chatbot(self): method to_openai_api_messages (line 315) | def to_openai_api_messages(self): method copy (line 327) | def copy(self): method dict (line 342) | def dict(self): function register_conv_template (line 355) | def register_conv_template(template: Conversation, override: bool = False): function get_conv_template (line 363) | def get_conv_template(name: str) -> Conversation: function calculate_ngram_repetition (line 388) | def calculate_ngram_repetition(text, n): function check_conversations_repetition (line 397) | def check_conversations_repetition(conversations, repeat_threshold=0.4, ... function get_frame_indices (line 406) | def get_frame_indices(num_frames, vlen, sample="rand", fix_start=None, i... function read_frames_gif (line 447) | def read_frames_gif(video_path, num_frames, sample="rand", fix_start=Non... function read_frames_decord (line 467) | def read_frames_decord(video_path, num_frames, sample="rand", fix_start=... function extract_frame_number (line 495) | def extract_frame_number(filename): function sort_frames (line 501) | def sort_frames(frame_paths): function read_frames_folder (line 506) | def read_frames_folder(video_path, num_frames, sample="rand", fix_start=... class WeightedConcatDataset (line 531) | class WeightedConcatDataset(ConcatDataset): method __init__ (line 532) | def __init__(self, datasets, weights): method __iter__ (line 538) | def __iter__(self): method __len__ (line 541) | def __len__(self): function pil_loader (line 545) | def pil_loader(img_str): class TCSLoader (line 551) | class TCSLoader(object): method __init__ (line 552) | def __init__(self, conf_path, sc_config_key="sensecore"): method __call__ (line 559) | def __call__(self, fn, image_type="image", max_num_frames=-1, min_num_... function expand2square (line 575) | def expand2square(pil_img, background_color): function simulate_jpeg_degradation (line 589) | def simulate_jpeg_degradation(quality): function build_transform (line 605) | def build_transform(is_train, input_size, pad2square=False, normalize_ty... function preprocess (line 629) | def preprocess(template_name, sources, tokenizer: transformers.PreTraine... function preprocess_mpt (line 721) | def preprocess_mpt(template_name, sources, tokenizer: transformers.PreTr... function preprocess_phi3 (line 802) | def preprocess_phi3(template_name, sources, tokenizer: transformers.PreT... function preprocess_internlm (line 898) | def preprocess_internlm(template_name, sources, tokenizer: transformers.... function preprocess_internvl2_5 (line 978) | def preprocess_internvl2_5(template_name, sources, tokenizer: transforme... function find_closest_aspect_ratio (line 1070) | def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height... function dynamic_preprocess (line 1087) | def dynamic_preprocess(image, min_num=1, max_num=6, image_size=448, use_... function preprocess_internvl2_5_siirl (line 1118) | def preprocess_internvl2_5_siirl(sources, tokenizer: transformers.PreTra... function find_closest_aspect_ratio (line 1205) | def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height... function dynamic_preprocess (line 1222) | def dynamic_preprocess(image, min_num=1, max_num=6, image_size=448, use_... FILE: siirl/models/transformers/internvl_chat/configuration_intern_vit.py class InternVisionConfig (line 16) | class InternVisionConfig(PretrainedConfig): method __init__ (line 64) | def __init__( method from_pretrained (line 108) | def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.... FILE: siirl/models/transformers/internvl_chat/configuration_internlm2.py class InternLM2Config (line 27) | class InternLM2Config(PretrainedConfig): method __init__ (line 78) | def __init__( # pylint: disable=W0102 method _rope_scaling_validation (line 132) | def _rope_scaling_validation(self): FILE: siirl/models/transformers/internvl_chat/configuration_internvl_chat.py class InternVLChatConfig (line 19) | class InternVLChatConfig(PretrainedConfig): method __init__ (line 23) | def __init__(self, vision_config=None, llm_config=None, use_backbone_l... method to_dict (line 58) | def to_dict(self): FILE: siirl/models/transformers/internvl_chat/modeling_intern_vit.py class FlashAttention (line 34) | class FlashAttention(nn.Module): method __init__ (line 45) | def __init__(self, softmax_scale=None, attention_dropout=0.0, device=N... method forward (line 50) | def forward(self, qkv, key_padding_mask=None, causal=False, cu_seqlens... class InternRMSNorm (line 85) | class InternRMSNorm(nn.Module): method __init__ (line 86) | def __init__(self, hidden_size, eps=1e-6): method forward (line 91) | def forward(self, hidden_states): class InternVisionEmbeddings (line 119) | class InternVisionEmbeddings(nn.Module): method __init__ (line 120) | def __init__(self, config: InternVisionConfig): method _get_pos_embed (line 138) | def _get_pos_embed(self, pos_embed, H, W): method forward (line 144) | def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor: class InternAttention (line 156) | class InternAttention(nn.Module): method __init__ (line 159) | def __init__(self, config: InternVisionConfig): method _naive_attn (line 186) | def _naive_attn(self, x): method _flash_attn (line 205) | def _flash_attn(self, x, key_padding_mask=None, need_weights=False): method forward (line 220) | def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class InternMLP (line 225) | class InternMLP(nn.Module): method __init__ (line 226) | def __init__(self, config: InternVisionConfig): method forward (line 233) | def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class InternVisionEncoderLayer (line 240) | class InternVisionEncoderLayer(nn.Module): method __init__ (line 241) | def __init__(self, config: InternVisionConfig, drop_path_rate: float): method forward (line 257) | def forward( class InternVisionEncoder (line 272) | class InternVisionEncoder(nn.Module): method __init__ (line 282) | def __init__(self, config: InternVisionConfig): method forward (line 290) | def forward( class InternVisionModel (line 331) | class InternVisionModel(PreTrainedModel): method __init__ (line 337) | def __init__(self, config: InternVisionConfig): method resize_pos_embeddings (line 344) | def resize_pos_embeddings(self, old_size, new_size, patch_size): method get_input_embeddings (line 356) | def get_input_embeddings(self): method forward (line 359) | def forward( FILE: siirl/models/transformers/internvl_chat/modeling_internlm2.py function _import_flash_attn (line 62) | def _import_flash_attn(): function _get_unpad_data (line 79) | def _get_unpad_data(attention_mask): function _make_causal_mask (line 92) | def _make_causal_mask(input_ids_shape: torch.Size, dtype: torch.dtype, d... function _expand_mask (line 108) | def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Option... class InternLM2RMSNorm (line 123) | class InternLM2RMSNorm(nn.Module): method __init__ (line 124) | def __init__(self, hidden_size, eps=1e-6): method forward (line 132) | def forward(self, hidden_states): class InternLM2RotaryEmbedding (line 141) | class InternLM2RotaryEmbedding(nn.Module): method __init__ (line 142) | def __init__(self, dim, max_position_embeddings=2048, base=10000, devi... method _set_cos_sin_cache (line 154) | def _set_cos_sin_cache(self, seq_len, device, dtype): method forward (line 164) | def forward(self, x, seq_len=None): class InternLM2LinearScalingRotaryEmbedding (line 176) | class InternLM2LinearScalingRotaryEmbedding(InternLM2RotaryEmbedding): method __init__ (line 179) | def __init__(self, dim, max_position_embeddings=2048, base=10000, devi... method _set_cos_sin_cache (line 183) | def _set_cos_sin_cache(self, seq_len, device, dtype): class InternLM2DynamicNTKScalingRotaryEmbedding (line 196) | class InternLM2DynamicNTKScalingRotaryEmbedding(InternLM2RotaryEmbedding): method __init__ (line 201) | def __init__(self, dim, max_position_embeddings=2048, base=10000, devi... method _set_cos_sin_cache (line 205) | def _set_cos_sin_cache(self, seq_len, device, dtype): function rotate_half (line 223) | def rotate_half(x): function apply_rotary_pos_emb (line 231) | def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1): class InternLM2MLP (line 240) | class InternLM2MLP(nn.Module): method __init__ (line 241) | def __init__(self, config): method forward (line 251) | def forward(self, x): function repeat_kv (line 258) | def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: class InternLM2Attention (line 271) | class InternLM2Attention(nn.Module): method __init__ (line 274) | def __init__(self, config: InternLM2Config): method _init_rope (line 297) | def _init_rope(self): method _shape (line 325) | def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): method forward (line 328) | def forward( class InternLM2FlashAttention2 (line 406) | class InternLM2FlashAttention2(InternLM2Attention): method forward (line 413) | def forward( method _flash_attention_forward (line 480) | def _flash_attention_forward(self, query_states, key_states, value_sta... method _unpad_input (line 528) | def _unpad_input(self, query_layer, key_layer, value_layer, attention_... class InternLM2DecoderLayer (line 567) | class InternLM2DecoderLayer(nn.Module): method __init__ (line 568) | def __init__(self, config: InternLM2Config): method forward (line 578) | def forward( class InternLM2PreTrainedModel (line 660) | class InternLM2PreTrainedModel(PreTrainedModel): method _init_weights (line 668) | def _init_weights(self, module): class InternLM2Model (line 750) | class InternLM2Model(InternLM2PreTrainedModel): method __init__ (line 760) | def __init__(self, config: InternLM2Config): method get_input_embeddings (line 778) | def get_input_embeddings(self): method set_input_embeddings (line 781) | def set_input_embeddings(self, value): method _prepare_decoder_attention_mask (line 784) | def _prepare_decoder_attention_mask(self, attention_mask, input_shape,... method forward (line 804) | def forward( class InternLM2ForCausalLM (line 928) | class InternLM2ForCausalLM(InternLM2PreTrainedModel): method __init__ (line 933) | def __init__(self, config): method get_input_embeddings (line 942) | def get_input_embeddings(self): method set_input_embeddings (line 945) | def set_input_embeddings(self, value): method get_output_embeddings (line 948) | def get_output_embeddings(self): method set_output_embeddings (line 951) | def set_output_embeddings(self, new_embeddings): method set_decoder (line 954) | def set_decoder(self, decoder): method get_decoder (line 957) | def get_decoder(self): method forward (line 962) | def forward( method prepare_inputs_for_generation (line 1050) | def prepare_inputs_for_generation(self, input_ids, past_key_values=Non... method _reorder_cache (line 1088) | def _reorder_cache(past_key_values, beam_idx): method build_inputs (line 1094) | def build_inputs(self, tokenizer, query: str, history: List[Tuple[str,... method chat (line 1107) | def chat( method stream_chat (line 1143) | def stream_chat( class InternLM2ForSequenceClassification (line 1242) | class InternLM2ForSequenceClassification(InternLM2PreTrainedModel): method __init__ (line 1243) | def __init__(self, config): method get_input_embeddings (line 1252) | def get_input_embeddings(self): method set_input_embeddings (line 1255) | def set_input_embeddings(self, value): method forward (line 1259) | def forward( FILE: siirl/models/transformers/internvl_chat/modeling_internvl_chat.py function version_cmp (line 28) | def version_cmp(v1, v2, op="eq"): class InternVLChatModel (line 37) | class InternVLChatModel(PreTrainedModel): method __init__ (line 45) | def __init__(self, config: InternVLChatConfig, vision_model=None, lang... method forward (line 86) | def forward( method pixel_shuffle (line 168) | def pixel_shuffle(self, x, scale_factor=0.5): method extract_feature (line 182) | def extract_feature(self, pixel_values): method batch_chat (line 196) | def batch_chat(self, tokenizer, pixel_values, questions, generation_co... method chat (line 238) | def chat(self, tokenizer, pixel_values, question, generation_config, h... method generate (line 287) | def generate( FILE: siirl/models/transformers/internvl_chat/tokenization_internlm2.py class InternLM2Tokenizer (line 35) | class InternLM2Tokenizer(PreTrainedTokenizer): method __init__ (line 49) | def __init__( method no_prefix_space_tokens (line 81) | def no_prefix_space_tokens(self): method vocab_size (line 88) | def vocab_size(self): method bos_token_id (line 93) | def bos_token_id(self) -> Optional[int]: method eos_token_id (line 97) | def eos_token_id(self) -> Optional[int]: method get_vocab (line 100) | def get_vocab(self): method _tokenize (line 106) | def _tokenize(self, text): method _convert_token_to_id (line 110) | def _convert_token_to_id(self, token): method _convert_id_to_token (line 114) | def _convert_id_to_token(self, index): method _maybe_add_prefix_space (line 119) | def _maybe_add_prefix_space(self, tokens, decoded): method convert_tokens_to_string (line 125) | def convert_tokens_to_string(self, tokens): method save_vocabulary (line 146) | def save_vocabulary(self, save_directory, filename_prefix: Optional[st... method build_inputs_with_special_tokens (line 171) | def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=No... method get_special_tokens_mask (line 187) | def get_special_tokens_mask(self, token_ids_0: List[int], token_ids_1:... method create_token_type_ids_from_sequences (line 210) | def create_token_type_ids_from_sequences(self, token_ids_0: List[int],... FILE: siirl/models/transformers/internvl_chat/tokenization_internlm2_fast.py class InternLM2Converter (line 37) | class InternLM2Converter(SpmConverter): method vocab (line 40) | def vocab(self, proto): method unk_id (line 49) | def unk_id(self, proto): method decoder (line 53) | def decoder(self, replacement, add_prefix_space): method tokenizer (line 63) | def tokenizer(self, proto): method normalizer (line 85) | def normalizer(self, proto): method pre_tokenizer (line 92) | def pre_tokenizer(self, replacement, add_prefix_space): class InternLM2TokenizerFast (line 100) | class InternLM2TokenizerFast(PreTrainedTokenizerFast): method __init__ (line 107) | def __init__( method can_save_slow_tokenizer (line 140) | def can_save_slow_tokenizer(self) -> bool: method update_post_processor (line 143) | def update_post_processor(self): method add_eos_token (line 168) | def add_eos_token(self): method add_bos_token (line 172) | def add_bos_token(self): method add_eos_token (line 176) | def add_eos_token(self, value): method add_bos_token (line 181) | def add_bos_token(self, value): method save_vocabulary (line 185) | def save_vocabulary(self, save_directory: str, filename_prefix: Option... FILE: siirl/models/transformers/kimi_vl.py function _merge_with_image_features (line 25) | def _merge_with_image_features( function rotate_half (line 70) | def rotate_half(x): function apply_rotary_pos_emb (line 78) | def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1): function repeat_kv (line 114) | def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: function _ulysses_flash_attn_forward (line 126) | def _ulysses_flash_attn_forward( FILE: siirl/models/transformers/llama.py function llama_flash_attn_forward (line 42) | def llama_flash_attn_forward( function llama_attn_forward (line 165) | def llama_attn_forward( class CausalLMOutputForPPO (line 236) | class CausalLMOutputForPPO(CausalLMOutputWithPast): function forward_for_ppo (line 241) | def forward_for_ppo( FILE: siirl/models/transformers/monkey_patch.py function repeat_kv (line 39) | def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: function _ulysses_flash_attention_forward (line 51) | def _ulysses_flash_attention_forward( function patch_vlm_for_ulysses_input_slicing (line 112) | def patch_vlm_for_ulysses_input_slicing(model_class: type): function apply_monkey_patch (line 143) | def apply_monkey_patch( function is_transformers_version_in_range (line 249) | def is_transformers_version_in_range(min_version: Optional[str] = None, ... FILE: siirl/models/transformers/npu_patch.py function apply_rotary_pos_emb_flashatt_npu (line 29) | def apply_rotary_pos_emb_flashatt_npu(q: torch.Tensor, k: torch.Tensor, ... function apply_rotary_pos_emb_npu (line 39) | def apply_rotary_pos_emb_npu(q, k, cos, sin, position_ids=None, unsqueez... function rms_norm_forward (line 48) | def rms_norm_forward(self, x): FILE: siirl/models/transformers/qwen2.py function qwen2_flash_attn_forward (line 33) | def qwen2_flash_attn_forward( function qwen2_attn_forward (line 149) | def qwen2_attn_forward( FILE: siirl/models/transformers/qwen2_5_vl.py class Qwen2_5_VLCausalLMOutputForPPO (line 26) | class Qwen2_5_VLCausalLMOutputForPPO(Qwen2_5_VLCausalLMOutputWithPast): function forward_for_ppo (line 31) | def forward_for_ppo( FILE: siirl/models/transformers/qwen2_vl.py function get_rope_index (line 54) | def get_rope_index( function prepare_fa2_from_position_ids (line 154) | def prepare_fa2_from_position_ids(query: torch.Tensor, key: torch.Tensor... function flash_attention_forward (line 170) | def flash_attention_forward( function ulysses_flash_attn_forward (line 233) | def ulysses_flash_attn_forward( class Qwen2VLCausalLMOutputForPPO (line 310) | class Qwen2VLCausalLMOutputForPPO(Qwen2VLCausalLMOutputWithPast): function forward_for_ppo (line 315) | def forward_for_ppo( FILE: siirl/models/transformers/transformers_compat.py function flash_attn_supports_top_left_mask (line 26) | def flash_attn_supports_top_left_mask(): FILE: siirl/models/weight_loader_registry.py function get_weight_loader (line 16) | def get_weight_loader(arch: str): function get_weight_saver (line 29) | def get_weight_saver(arch: str): FILE: siirl/params/dag_args.py class DagArguments (line 20) | class DagArguments: FILE: siirl/params/data_args.py class DataArguments (line 23) | class DataArguments: method __post_init__ (line 128) | def __post_init__(self): method to_dict (line 149) | def to_dict(self) -> Dict[str, Any]: FILE: siirl/params/display_dict.py function _render_dict_recursively_util (line 27) | def _render_dict_recursively_util(current_dict_to_render: Dict[str, Any]... function log_dict_formatted (line 71) | def log_dict_formatted(data_dict: Dict[str, Any], title: Optional[str] =... FILE: siirl/params/embodied_args.py class EnvironmentArgs (line 20) | class EnvironmentArgs: class EmbodiedArguments (line 53) | class EmbodiedArguments: method to_dict (line 117) | def to_dict(self) -> Dict[str, Any]: class EmbodiedSamplingConfig (line 121) | class EmbodiedSamplingConfig: method to_dict (line 145) | def to_dict(self) -> Dict[str, Any]: FILE: siirl/params/model_args.py class MixedPrecisionArguments (line 22) | class MixedPrecisionArguments: class FSDPArguments (line 43) | class FSDPArguments: method to_dict (line 60) | def to_dict(self) -> Dict[str, Any]: class MegatronArguments (line 65) | class MegatronArguments: method to_dict (line 93) | def to_dict(self) -> Dict[str, Any]: class OptimizerArguments (line 98) | class OptimizerArguments: method to_dict (line 124) | def to_dict(self) -> Dict[str, Any]: class ProcessorArguments (line 129) | class ProcessorArguments: class ModelArguments (line 161) | class ModelArguments(ProcessorArguments): method __post_init__ (line 228) | def __post_init__(self): method to_dict (line 238) | def to_dict(self) -> Dict[str, Any]: class CheckpointArguments (line 243) | class CheckpointArguments: class PolicyLossArguments (line 260) | class PolicyLossArguments: class ActorArguments (line 274) | class ActorArguments: method to_dict (line 339) | def to_dict(self) -> Dict[str, Any]: class EvalSamplingArguments (line 344) | class EvalSamplingArguments: class LayerNameMapArguments (line 353) | class LayerNameMapArguments: method to_dict (line 357) | def to_dict(self) -> Dict[str, Any]: class MultiTurnArguments (line 362) | class MultiTurnArguments: class CustomAsyncServer (line 410) | class CustomAsyncServer: class AgentArguments (line 418) | class AgentArguments: class EngineArguments (line 436) | class EngineArguments: class RolloutArguments (line 442) | class RolloutArguments: method to_dict (line 493) | def to_dict(self) -> Dict[str, Any]: class RefArguments (line 498) | class RefArguments: method to_dict (line 533) | def to_dict(self) -> Dict[str, Any]: class ActorRolloutRefArguments (line 538) | class ActorRolloutRefArguments: method to_dict (line 546) | def to_dict(self) -> Dict[str, Any]: class CriticArguments (line 551) | class CriticArguments: method to_dict (line 583) | def to_dict(self) -> Dict[str, Any]: class OverlongBufferArguments (line 588) | class OverlongBufferArguments: method to_dict (line 596) | def to_dict(self) -> Dict[str, Any]: class RewardModelArguments (line 601) | class RewardModelArguments: method to_dict (line 635) | def to_dict(self) -> Dict[str, Any]: class KLCtrlArguments (line 640) | class KLCtrlArguments: class FilterGroupsArguments (line 648) | class FilterGroupsArguments: method to_dict (line 657) | def to_dict(self) -> Dict[str, Any]: class AlgorithmArguments (line 662) | class AlgorithmArguments: method to_dict (line 689) | def to_dict(self) -> Dict[str, Any]: FILE: siirl/params/parser.py function _set_transformers_logging (line 26) | def _set_transformers_logging() -> None: function parse_config (line 33) | def parse_config() -> SiiRLArguments: FILE: siirl/params/profiler_args.py class ProfilerArguments (line 20) | class ProfilerArguments: method to_dict (line 39) | def to_dict(self) -> Dict[str, Any]: FILE: siirl/params/training_args.py class TrainingArguments (line 29) | class TrainingArguments: method to_dict (line 78) | def to_dict(self) -> Dict[str, Any]: class CustomRewardArguments (line 83) | class CustomRewardArguments: class SiiRLArguments (line 90) | class SiiRLArguments: method to_dict (line 101) | def to_dict(self) -> Dict[str, Any]: FILE: siirl/third_party/sglang/parallel_state.py function initialize_parallel_state (line 38) | def initialize_parallel_state( function ensure_model_parallel_initialized (line 77) | def ensure_model_parallel_initialized( function model_parallel_is_initialized (line 105) | def model_parallel_is_initialized(): function initialize_model_parallel_for_sglang (line 111) | def initialize_model_parallel_for_sglang( function initialize_model_parallel (line 204) | def initialize_model_parallel( function get_device_mesh (line 292) | def get_device_mesh(): function get_tensor_model_parallel_group (line 306) | def get_tensor_model_parallel_group(): function get_tensor_model_parallel_world_size (line 313) | def get_tensor_model_parallel_world_size(): function get_tensor_model_parallel_rank (line 318) | def get_tensor_model_parallel_rank(): function get_tensor_model_parallel_src_rank (line 323) | def get_tensor_model_parallel_src_rank(): FILE: siirl/user_interface/filter_interface/dapo.py function dynamic_sampling (line 24) | def dynamic_sampling(config: SiiRLArguments, batch: TensorDict, **kwargs... FILE: siirl/user_interface/filter_interface/embodied.py function verify (line 28) | def verify( function _filter_batch (line 74) | def _filter_batch(batch: TensorDict, n_samples: int, config: SiiRLArgume... function _compute_embodied_verification_metrics (line 145) | def _compute_embodied_verification_metrics( function embodied_local_rank_sampling (line 192) | def embodied_local_rank_sampling( FILE: siirl/user_interface/rewards_interface/custom_gsm8k_reward.py function extract_solution (line 18) | def extract_solution(solution_str, method="strict"): function compute_score (line 43) | def compute_score(data_source, solution_str, ground_truth, extra_info): FILE: siirl/utils/checkpoint/checkpoint_manager.py class BaseCheckpointManager (line 33) | class BaseCheckpointManager: method __init__ (line 48) | def __init__( method should_save_model (line 88) | def should_save_model(self) -> bool: method should_save_optimizer (line 95) | def should_save_optimizer(self) -> bool: method should_save_extra (line 102) | def should_save_extra(self) -> bool: method should_save_hf_model (line 109) | def should_save_hf_model(self) -> bool: method should_load_model (line 117) | def should_load_model(self) -> bool: method should_load_optimizer (line 124) | def should_load_optimizer(self) -> bool: method should_load_extra (line 131) | def should_load_extra(self) -> bool: method load_checkpoint (line 137) | def load_checkpoint(self, local_path: str, hdfs_path: str = None, del_... method save_checkpoint (line 140) | def save_checkpoint( method checkpath (line 146) | def checkpath(local_path: str, hdfs_path: str): method remove_previous_save_local_path (line 150) | def remove_previous_save_local_path(self, path): method local_mkdir (line 165) | def local_mkdir(path): method get_rng_state (line 186) | def get_rng_state(): method load_rng_state (line 201) | def load_rng_state(rng_state): function find_latest_ckpt_path (line 212) | def find_latest_ckpt_path(path, directory_format="global_step_{}"): function get_checkpoint_tracker_filename (line 244) | def get_checkpoint_tracker_filename(root_path: str): FILE: siirl/utils/checkpoint/fsdp_checkpoint_manager.py class FSDPCheckpointManager (line 33) | class FSDPCheckpointManager(BaseCheckpointManager): method __init__ (line 51) | def __init__( method load_checkpoint (line 76) | def load_checkpoint(self, local_path: str, hdfs_path: str = None, del_... method save_checkpoint (line 129) | def save_checkpoint(self, local_path: str, hdfs_path: str = None, glob... FILE: siirl/utils/checkpoint/megatron_checkpoint_manager.py class MegatronCheckpointManager (line 44) | class MegatronCheckpointManager(BaseCheckpointManager): method __init__ (line 98) | def __init__( method get_rng_state (line 147) | def get_rng_state(self, use_dist_ckpt: bool = True, data_parallel_rand... method get_checkpoint_name (line 181) | def get_checkpoint_name( method generate_state_dict (line 227) | def generate_state_dict( method load_rng_states (line 267) | def load_rng_states(self, rng_states, data_parallel_random_init=False,... method load_checkpoint (line 285) | def load_checkpoint(self, local_path: str, hdfs_path: str = None, del_... method save_checkpoint (line 356) | def save_checkpoint(self, local_path: str, hdfs_path: str = None, glob... FILE: siirl/utils/debug/mstx_profile.py function mark_start_range (line 30) | def mark_start_range(message: Optional[str] = None) -> None: function mark_end_range (line 40) | def mark_end_range(range_id: str) -> None: function mark_annotate (line 50) | def mark_annotate(message: Optional[str] = None) -> Callable: function marked_timer (line 66) | def marked_timer(name: str, timing_raw: dict[str, float], *args: Any, **... function get_npu_profiler (line 90) | def get_npu_profiler(config: ProfilerArguments, role: Optional[str] = No... class NPUProfiler (line 136) | class NPUProfiler(DistProfiler): method __init__ (line 143) | def __init__(self, rank: int, config: ProfilerArguments, **kwargs): method start (line 149) | def start(self, **kwargs): method stop (line 159) | def stop(self): method annotate (line 168) | def annotate(message: Optional[str] = None, role: Optional[str] = None... FILE: siirl/utils/debug/performance.py function _get_current_mem_info (line 26) | def _get_current_mem_info(unit: str = "GB", precision: int = 2) -> Tuple... function log_gpu_memory_usage (line 44) | def log_gpu_memory_usage(head: str, logger=None, level="DEBUG", rank: in... class GPUMemoryLogger (line 55) | class GPUMemoryLogger(DecoratorLoggerBase): method __init__ (line 65) | def __init__(self, role: str, logger: logging.Logger = None, level=log... method __call__ (line 72) | def __call__(self, decorated_function: callable): method log (line 78) | def log(self, func, *args, **kwargs): function log_print (line 93) | def log_print(ctn: Any): FILE: siirl/utils/debug/profile.py class Profiler (line 26) | class Profiler: method __init__ (line 27) | def __init__(self, config): method _validate (line 54) | def _validate(self): method check (line 63) | def check(self): method start (line 66) | def start(self): method step (line 71) | def step(self): method stop (line 75) | def stop(self): method save (line 80) | def save(self): method stop_and_save (line 90) | def stop_and_save(self): method stop_trace (line 95) | def stop_trace(self): function mark_start_range (line 101) | def mark_start_range( function mark_end_range (line 118) | def mark_end_range(range_id: str) -> None: function mark_annotate (line 127) | def mark_annotate( class DistProfiler (line 151) | class DistProfiler: method __init__ (line 153) | def __init__(self, rank: int, config: ProfilerArguments, **kwargs): method start (line 159) | def start(self, **kwargs): method stop (line 162) | def stop(self): method annotate (line 166) | def annotate( FILE: siirl/utils/embodied/libero_utils.py function get_libero_env (line 18) | def get_libero_env(task, model_family, gpu_id=-1, resolution=256): function get_libero_dummy_action (line 33) | def get_libero_dummy_action(model_family: str): function resize_image (line 38) | def resize_image(img, resize_size): function get_libero_image (line 56) | def get_libero_image(obs, resize_size): function get_libero_wrist_image (line 67) | def get_libero_wrist_image(obs, resize_size): function quat2axisangle (line 93) | def quat2axisangle(quat): function get_image_resize_size (line 119) | def get_image_resize_size(cfg): function normalize_gripper_action (line 150) | def normalize_gripper_action(action: np.ndarray, binarize: bool = True) ... function invert_gripper_action (line 190) | def invert_gripper_action(action: np.ndarray) -> np.ndarray: function save_rollout_video (line 211) | def save_rollout_video(rollout_images, exp_name, task_name, step_idx, su... FILE: siirl/utils/embodied/openvla_utils.py function update_auto_map (line 27) | def update_auto_map(pretrained_checkpoint: str) -> None: function check_identical_files (line 128) | def check_identical_files(path1: Union[str, Path], path2: Union[str, Pat... function _handle_file_sync (line 149) | def _handle_file_sync(curr_filepath: str, checkpoint_filepath: str, file... function check_model_logic_mismatch (line 197) | def check_model_logic_mismatch(pretrained_checkpoint: str) -> None: FILE: siirl/utils/embodied/video_emb.py class VideoEmbeddingModel (line 29) | class VideoEmbeddingModel: method __init__ (line 34) | def __init__(self, model_path: str, img_size: int = 384, device_id: in... method _build_pt_video_transform (line 46) | def _build_pt_video_transform(self): method _load_pretrained_vjepa_pt_weights (line 59) | def _load_pretrained_vjepa_pt_weights(self, model, pretrained_weights): method _create_model_instance (line 71) | def _create_model_instance(self): method offload_to_host (line 81) | def offload_to_host(self): method load_to_device (line 87) | def load_to_device(self): method extract_video_embedding (line 92) | def extract_video_embedding(self, video_tensor): method extract_video_embedding_batch (line 109) | def extract_video_embedding_batch(self, video_tensor_list): method get_embeddings (line 124) | def get_embeddings(self, batch_names, batch_frames): FILE: siirl/utils/experimental/torch_functional.py function _fused_linear_for_ppo_fwd (line 20) | def _fused_linear_for_ppo_fwd(hidden_states: torch.FloatTensor, vocab_we... function _fused_linear_for_ppo_bwd (line 35) | def _fused_linear_for_ppo_bwd( class FusedLinearForPPOFunction (line 70) | class FusedLinearForPPOFunction(torch.autograd.Function): method forward (line 72) | def forward( method backward (line 127) | def backward(ctx, dlog_probs: Optional[torch.FloatTensor], dentropy: O... class FusedLinearForPPO (line 191) | class FusedLinearForPPO(torch.nn.Module): method __init__ (line 192) | def __init__(self, chunk_size: int = 512): method forward (line 197) | def forward( FILE: siirl/utils/extras/device.py function is_torch_npu_available (line 18) | def is_torch_npu_available() -> bool: function get_device_name (line 32) | def get_device_name() -> str: function get_torch_device (line 47) | def get_torch_device() -> any: function get_device_id (line 60) | def get_device_id() -> int: function get_nccl_backend (line 67) | def get_nccl_backend() -> str: function device_synchronize (line 80) | def device_synchronize(): function set_expandable_segments (line 94) | def set_expandable_segments(enable: bool) -> None: FILE: siirl/utils/extras/fs.py function is_non_local (line 34) | def is_non_local(path): function md5_encode (line 46) | def md5_encode(path: str) -> str: function get_local_temp_path (line 61) | def get_local_temp_path(hdfs_path: str, cache_dir: str) -> str: function verify_copy (line 82) | def verify_copy(src: str, dest: str) -> bool: function copy_to_shm (line 141) | def copy_to_shm(src: str): function _record_directory_structure (line 161) | def _record_directory_structure(folder_path): function _check_directory_structure (line 175) | def _check_directory_structure(folder_path, record_file): function copy_to_local (line 192) | def copy_to_local(src: str, cache_dir=None, filelock=".file.lock", verbo... function copy_local_path_from_hdfs (line 214) | def copy_local_path_from_hdfs(src: str, cache_dir=None, filelock=".file.... function local_mkdir_safe (line 256) | def local_mkdir_safe(path): FILE: siirl/utils/extras/hdfs_io.py function exists (line 25) | def exists(path: str, **kwargs) -> bool: function _exists (line 41) | def _exists(file_path: str): function makedirs (line 48) | def makedirs(name, mode=0o777, exist_ok=False, **kwargs) -> None: function _mkdir (line 73) | def _mkdir(file_path: str) -> bool: function copy (line 82) | def copy(src: str, dst: str, **kwargs) -> bool: function _copy (line 111) | def _copy(from_path: str, to_path: str, timeout: int = None) -> bool: function _run_cmd (line 138) | def _run_cmd(cmd: str, timeout=None): function _hdfs_cmd (line 142) | def _hdfs_cmd(cmd: str) -> str: function _is_non_local (line 146) | def _is_non_local(path: str): FILE: siirl/utils/extras/import_utils.py function is_megatron_core_available (line 25) | def is_megatron_core_available(): function is_vllm_available (line 34) | def is_vllm_available(): function is_sglang_available (line 43) | def is_sglang_available(): function is_nvtx_available (line 52) | def is_nvtx_available(): function import_external_libs (line 60) | def import_external_libs(external_libs=None): function load_extern_type (line 71) | def load_extern_type(file_path: Optional[str], type_name: Optional[str]): function _get_qualified_name (line 95) | def _get_qualified_name(func): function deprecated (line 102) | def deprecated(replacement: str = ""): FILE: siirl/utils/extras/misc.py class AverageMeter (line 35) | class AverageMeter: method __init__ (line 40) | def __init__(self): method reset (line 43) | def reset(self): method update (line 49) | def update(self, val, n=1): function check_version (line 56) | def check_version(requirement: str, mandatory: bool = False) -> None: function check_dependencies (line 72) | def check_dependencies() -> None: function calculate_tps (line 85) | def calculate_tps( function count_parameters (line 104) | def count_parameters(model: "torch.nn.Module") -> Tuple[int, int]: function get_current_device (line 133) | def get_current_device() -> "torch.device": function get_device_count (line 151) | def get_device_count() -> int: function get_logits_processor (line 165) | def get_logits_processor() -> "LogitsProcessorList": function get_peak_memory (line 174) | def get_peak_memory() -> Tuple[int, int]: function has_tokenized_data (line 186) | def has_tokenized_data(path: "os.PathLike") -> bool: function infer_optim_dtype (line 193) | def infer_optim_dtype(model_dtype: "torch.dtype") -> "torch.dtype": function is_gpu_or_npu_available (line 205) | def is_gpu_or_npu_available() -> bool: function is_env_enabled (line 212) | def is_env_enabled(env_var: str, default: str = "0") -> bool: function numpify (line 219) | def numpify(inputs: Union["NDArray", "torch.Tensor"]) -> "NDArray": function skip_check_imports (line 233) | def skip_check_imports() -> None: function torch_gc (line 241) | def torch_gc() -> None: function try_download_model_from_other_hub (line 256) | def try_download_model_from_other_hub(model_args: "ModelArguments") -> str: function use_modelscope (line 282) | def use_modelscope() -> bool: function use_openmind (line 286) | def use_openmind() -> bool: function use_ray (line 290) | def use_ray() -> bool: FILE: siirl/utils/extras/net_utils.py function is_ipv4 (line 30) | def is_ipv4(ip_str: str) -> bool: function is_ipv6 (line 47) | def is_ipv6(ip_str: str) -> bool: FILE: siirl/utils/extras/packages.py function _is_package_available (line 12) | def _is_package_available(name: str) -> bool: function _get_package_version (line 16) | def _get_package_version(name: str) -> "Version": function is_pyav_available (line 23) | def is_pyav_available(): function is_librosa_available (line 27) | def is_librosa_available(): function is_fastapi_available (line 31) | def is_fastapi_available(): function is_galore_available (line 35) | def is_galore_available(): function is_apollo_available (line 39) | def is_apollo_available(): function is_gradio_available (line 43) | def is_gradio_available(): function is_matplotlib_available (line 47) | def is_matplotlib_available(): function is_pillow_available (line 51) | def is_pillow_available(): function is_ray_available (line 55) | def is_ray_available(): function is_requests_available (line 59) | def is_requests_available(): function is_rouge_available (line 63) | def is_rouge_available(): function is_starlette_available (line 67) | def is_starlette_available(): function is_transformers_version_greater_than (line 72) | def is_transformers_version_greater_than(content: str): function is_uvicorn_available (line 76) | def is_uvicorn_available(): function is_vllm_available (line 80) | def is_vllm_available(): FILE: siirl/utils/extras/patch.py function verify (line 24) | def verify( FILE: siirl/utils/extras/py_functional.py function _mp_target_wrapper (line 30) | def _mp_target_wrapper(target_func: Callable, mp_queue: multiprocessing.... function timeout_limit (line 51) | def timeout_limit(seconds: float, use_signals: bool = False): function union_two_dict (line 141) | def union_two_dict(dict1: Dict, dict2: Dict): function append_to_dict (line 159) | def append_to_dict(data: Dict, new_data: Dict): class NestedNamespace (line 178) | class NestedNamespace(SimpleNamespace): method __init__ (line 194) | def __init__(self, dictionary, **kwargs): class DynamicEnumMeta (line 203) | class DynamicEnumMeta(type): method __iter__ (line 204) | def __iter__(cls) -> Iterator[Any]: method __contains__ (line 207) | def __contains__(cls, item: Any) -> bool: method __getitem__ (line 213) | def __getitem__(cls, name: str) -> Any: method __reduce_ex__ (line 216) | def __reduce_ex__(cls, protocol): method names (line 220) | def names(cls): method values (line 223) | def values(cls): class DynamicEnum (line 227) | class DynamicEnum(metaclass=DynamicEnumMeta): method __init__ (line 231) | def __init__(self, name: str, value: int): method __repr__ (line 235) | def __repr__(self): method __reduce_ex__ (line 238) | def __reduce_ex__(self, protocol): method register (line 248) | def register(cls, name: str) -> "DynamicEnum": method remove (line 259) | def remove(cls, name: str): method from_name (line 266) | def from_name(cls, name: str) -> Optional["DynamicEnum"]: function convert_to_regular_types (line 270) | def convert_to_regular_types(obj): FILE: siirl/utils/extras/ray_utils.py function ray_noset_visible_devices (line 25) | def ray_noset_visible_devices(env_vars=os.environ): function parallel_put (line 48) | def parallel_put(data_list: List[Any], max_workers: Optional[int] = None): FILE: siirl/utils/import_string.py function import_string (line 18) | def import_string(import_name: str): FILE: siirl/utils/kernel/kernels.py function null_decorator (line 56) | def null_decorator(*args, **kwargs): class EntropyReductionEnum (line 73) | class EntropyReductionEnum: function get_entropy_reduction_enum_number (line 83) | def get_entropy_reduction_enum_number(reduction: str) -> int: function get_entropy_reduction_enum (line 99) | def get_entropy_reduction_enum(ce_reduction: int) -> EntropyReductionEnum: class BackwardEnum (line 116) | class BackwardEnum: class Config (line 130) | class Config: function set_backward_method (line 145) | def set_backward_method(backward_method: BackwardEnum): function efficient_entropy_kernel_general_mainloop (line 158) | def efficient_entropy_kernel_general_mainloop( function efficient_entropy_triton_kernel_epilogue (line 292) | def efficient_entropy_triton_kernel_epilogue( function efficient_entropy_triton_kernel_epilogue_tp (line 384) | def efficient_entropy_triton_kernel_epilogue_tp( function efficient_entropy_triton_epilogue_tp_update (line 460) | def efficient_entropy_triton_epilogue_tp_update( function efficient_entropy_forward (line 507) | def efficient_entropy_forward( function efficient_entropy_backward_kernel_general_mainloop_MN (line 711) | def efficient_entropy_backward_kernel_general_mainloop_MN( function efficient_entropy_backward_kernel_d_hidden (line 888) | def efficient_entropy_backward_kernel_d_hidden( function efficient_entropy_backward_kernel_d_weight (line 1015) | def efficient_entropy_backward_kernel_d_weight( function efficient_entropy_backward_kernel_general_d_logits (line 1135) | def efficient_entropy_backward_kernel_general_d_logits( function efficient_entropy_backward_kernel_general_d_logits_split_N (line 1274) | def efficient_entropy_backward_kernel_general_d_logits_split_N( function efficient_entropy_backward (line 1378) | def efficient_entropy_backward( FILE: siirl/utils/kernel/linear_cross_entropy.py class LinearCrossEntropy (line 40) | class LinearCrossEntropy(torch.autograd.Function): method forward (line 42) | def forward( method backward (line 90) | def backward(ctx, dlogprobs: torch.Tensor, dentropy: torch.Tensor) -> ... FILE: siirl/utils/logger/aggregate_logger.py function concat_dict_to_str (line 24) | def concat_dict_to_str(dict: Dict, step): class LocalLogger (line 33) | class LocalLogger: method __init__ (line 34) | def __init__(self, remote_logger=None, enable_wandb=False, print_to_co... method flush (line 37) | def flush(self): method log (line 40) | def log(self, data, step): class DecoratorLoggerBase (line 45) | class DecoratorLoggerBase: method __init__ (line 46) | def __init__(self, role: str, logger: logging.Logger = None, level=log... method log_by_print (line 56) | def log_by_print(self, log_str): method log_by_logging (line 60) | def log_by_logging(self, log_str): function log_with_rank (line 67) | def log_with_rank(message: str, rank, logger: logging.Logger, level=logg... FILE: siirl/utils/logger/logging_utils.py function set_basic_config (line 26) | def set_basic_config(): FILE: siirl/utils/logger/tracking.py class Tracking (line 27) | class Tracking: method __init__ (line 40) | def __init__(self, project_name, experiment_name, default_backend: Uni... method log (line 132) | def log(self, data, step, backend=None): method __del__ (line 137) | def __del__(self): class ClearMLLogger (line 151) | class ClearMLLogger: method __init__ (line 152) | def __init__(self, project_name: str, experiment_name: str, config): method _get_logger (line 167) | def _get_logger(self): method log (line 170) | def log(self, data, step): method finish (line 196) | def finish(self): class _TensorboardAdapter (line 200) | class _TensorboardAdapter: method __init__ (line 201) | def __init__(self): method log (line 211) | def log(self, data, step): method finish (line 215) | def finish(self): class _MlflowLoggingAdapter (line 219) | class _MlflowLoggingAdapter: method log (line 220) | def log(self, data, step): function _compute_mlflow_params_from_objects (line 227) | def _compute_mlflow_params_from_objects(params) -> Dict[str, Any]: function _transform_params_to_json_serializable (line 234) | def _transform_params_to_json_serializable(x, convert_list_to_dict: bool): function _flatten_dict (line 254) | def _flatten_dict(raw: Dict[str, Any], *, sep: str) -> Dict[str, Any]: class ValidationGenerationsLogger (line 263) | class ValidationGenerationsLogger: method log (line 264) | def log(self, loggers, samples, step): method log_generations_to_wandb (line 277) | def log_generations_to_wandb(self, samples, step): method log_generations_to_swanlab (line 304) | def log_generations_to_swanlab(self, samples, step): method log_generations_to_mlflow (line 326) | def log_generations_to_mlflow(self, samples, step): method log_generations_to_clearml (line 348) | def log_generations_to_clearml(self, samples, step): method log_generations_to_tensorboard (line 376) | def log_generations_to_tensorboard(self, samples, step): FILE: siirl/utils/megatron/dist_checkpointing.py function save_dist_checkpointing (line 32) | def save_dist_checkpointing(sharded_state_dict, ckpt_path, async_save=Fa... function load_dist_checkpointing (line 52) | def load_dist_checkpointing(sharded_state_dict, ckpt_dir): FILE: siirl/utils/megatron/megatron_utils.py function get_model_config (line 45) | def get_model_config(model): function get_model (line 49) | def get_model( class McoreModuleWrapperConfig (line 162) | class McoreModuleWrapperConfig: function make_megatron_module (line 171) | def make_megatron_module( function unwrap_model (line 222) | def unwrap_model(model, module_instances=ALL_MODULE_WRAPPER_CLASSNAMES): function convert_config (line 237) | def convert_config(hf_config: PretrainedConfig, megatron_config) -> Tran... function init_megatron_optim_config (line 279) | def init_megatron_optim_config(optim_config: Dict) -> OptimizerConfig: function mcore_model_parallel_config (line 302) | def mcore_model_parallel_config( function offload_megatron_model_to_cpu (line 328) | def offload_megatron_model_to_cpu(models): function load_megatron_model_to_gpu (line 364) | def load_megatron_model_to_gpu(models, load_grad=True): function offload_megatron_copy_params (line 393) | def offload_megatron_copy_params(optimizers): function load_megatron_copy_params (line 434) | def load_megatron_copy_params(optimizers): function offload_megatron_optimizer (line 475) | def offload_megatron_optimizer(optimizers): function load_megatron_optimizer (line 494) | def load_megatron_optimizer(optimizers): function print_rank_0 (line 512) | def print_rank_0(message): function get_dist_checkpoint_path (line 521) | def get_dist_checkpoint_path(checkpoint_path): function get_hf_model_checkpoint_path (line 527) | def get_hf_model_checkpoint_path(checkpoint_path): function get_transformer_config_checkpoint_path (line 533) | def get_transformer_config_checkpoint_path(checkpoint_path): function get_model_checkpoint_path (line 538) | def get_model_checkpoint_path(checkpoint_path): function get_hf_config_and_tokenizer_checkpoint_path (line 543) | def get_hf_config_and_tokenizer_checkpoint_path(checkpoint_path): function get_optimizer_checkpoint_path (line 548) | def get_optimizer_checkpoint_path(checkpoint_path, use_distributed_optim... function get_rng_states_checkpoint_path (line 560) | def get_rng_states_checkpoint_path(checkpoint_path, only_rank0_save=True): function convert_megatron_model_to_transformers_model (line 572) | def convert_megatron_model_to_transformers_model( function broadcast_from_megatron_pp (line 712) | def broadcast_from_megatron_pp(tensor: torch.Tensor): function broadcast_str_from_megatron_pp (line 747) | def broadcast_str_from_megatron_pp(obj: Any): function default_tp_concat_fn (line 771) | def default_tp_concat_fn(layer_name_mapping, name, train_params, infer_p... function per_tensor_generator (line 829) | def per_tensor_generator(actor_module, model_config, weight_converter, t... function get_transformer_layer_offset (line 947) | def get_transformer_layer_offset(pipeline_rank, vp_rank, config: Transfo... FILE: siirl/utils/megatron/memory.py class MemoryBuffer (line 19) | class MemoryBuffer: method __init__ (line 20) | def __init__(self, numel, numel_padded, dtype): method zero (line 26) | def zero(self): method get (line 30) | def get(self, shape, start_index): FILE: siirl/utils/megatron/memory_buffer.py class MemoryBuffer (line 26) | class MemoryBuffer: method __init__ (line 32) | def __init__(self, numel: int, numel_padded: int, dtype: torch.dtype, ... method zero (line 41) | def zero(self): method get (line 45) | def get(self, shape, start_index): function calc_padded_numel (line 55) | def calc_padded_numel(shape: torch.Size, dtype: torch.dtype): function get_weight_buffer_meta_from_module (line 62) | def get_weight_buffer_meta_from_module(module: nn.Module) -> Dict[str, D... function build_memory_buffer (line 72) | def build_memory_buffer(weight_buffer_meta: Dict[str, Dict]) -> Dict[tor... function build_memory_reference_from_module (line 101) | def build_memory_reference_from_module(module: torch.nn.Module, memory_b... function build_memory_reference (line 115) | def build_memory_reference(weight_buffer_meta: Dict[str, Dict], memory_b... class MemoryBufferModuleWrapper (line 142) | class MemoryBufferModuleWrapper: method __init__ (line 148) | def __init__(self, module: nn.Module): method get_memory_buffers (line 155) | def get_memory_buffers(self): method get_weight_buffer_meta (line 158) | def get_weight_buffer_meta(self): class MegatronMemoryBufferForRollout (line 162) | class MegatronMemoryBufferForRollout: method __init__ (line 177) | def __init__(self, transform_memory_param_fn): method initialize_weight_buffer (line 183) | def initialize_weight_buffer(self, weight_buffer_meta_pp: List[Dict[st... method build_memory_reference (line 201) | def build_memory_reference(self): method named_parameters (line 207) | def named_parameters(self): method weight_buffers (line 211) | def weight_buffers(self): method memory_buffers (line 215) | def memory_buffers(self): FILE: siirl/utils/megatron/optimizer.py function get_megatron_optimizer (line 21) | def get_megatron_optimizer( function get_megatron_optimizer_param_scheduler (line 38) | def get_megatron_optimizer_param_scheduler( function get_megatron_last_lr (line 78) | def get_megatron_last_lr(optimizer): FILE: siirl/utils/megatron/pipeline_parallel.py function compute_transformers_input_shapes (line 22) | def compute_transformers_input_shapes(batches, meta_info): function make_batch_generator (line 49) | def make_batch_generator(batches, vpp_size): FILE: siirl/utils/megatron/sequence_parallel.py function mark_parameter_as_sequence_parallel (line 21) | def mark_parameter_as_sequence_parallel(parameter): function is_sequence_parallel_param (line 25) | def is_sequence_parallel_param(param): function pad_to_sequence_parallel (line 29) | def pad_to_sequence_parallel(unpad_tokens: torch.Tensor): FILE: siirl/utils/megatron/tensor_parallel.py function update_kwargs_with_config (line 30) | def update_kwargs_with_config(dictionary: Dict, config: "ModelParallelCo... function get_default_kwargs_for_model_parallel_config (line 35) | def get_default_kwargs_for_model_parallel_config(): function get_default_model_parallel_config (line 46) | def get_default_model_parallel_config(): function get_common_default_kwargs_for_parallel_linear (line 52) | def get_common_default_kwargs_for_parallel_linear(): function get_default_kwargs_for_column_parallel_linear (line 63) | def get_default_kwargs_for_column_parallel_linear(): function get_default_kwargs_for_row_parallel_linear (line 79) | def get_default_kwargs_for_row_parallel_linear(): function get_default_kwargs_for_parallel_embedding (line 84) | def get_default_kwargs_for_parallel_embedding(): function is_tensor_parallel_param (line 95) | def is_tensor_parallel_param(param): function get_tensor_parallel_partition_dim (line 99) | def get_tensor_parallel_partition_dim(param): function get_tensor_parallel_partition_stride (line 104) | def get_tensor_parallel_partition_stride(param): class _VocabParallelEntropy (line 109) | class _VocabParallelEntropy(torch.autograd.Function): method forward (line 111) | def forward(ctx, vocab_parallel_logits: torch.Tensor) -> torch.Tensor: method backward (line 130) | def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor: function vocab_parallel_entropy (line 142) | def vocab_parallel_entropy(vocab_parallel_logits: torch.Tensor) -> torch... function vocab_parallel_log_probs_from_logits (line 154) | def vocab_parallel_log_probs_from_logits(logits, labels): function vocab_parallel_log_probs_from_logits_response_rmpad (line 161) | def vocab_parallel_log_probs_from_logits_response_rmpad(input_ids, atten... FILE: siirl/utils/memory_utils.py function aggressive_empty_cache (line 24) | def aggressive_empty_cache(force_sync: bool = True, max_retries: int = 3... function reset_memory_stats (line 69) | def reset_memory_stats() -> None: function get_memory_info (line 77) | def get_memory_info() -> dict: function log_memory_usage (line 95) | def log_memory_usage(stage: str = "current") -> None: function optimize_memory_for_inference (line 110) | def optimize_memory_for_inference() -> None: function optimize_memory_for_training (line 124) | def optimize_memory_for_training() -> None: FILE: siirl/utils/metrics/metric_utils.py function _compute_response_info (line 37) | def _compute_response_info(batch: TensorDict) -> Dict[str, Any]: function compute_data_metric (line 70) | def compute_data_metric(data: TensorDict): function compute_timing_metrics (line 186) | def compute_timing_metrics(batch: TensorDict, timing_raw: Dict[str, floa... function compute_throughout_metrics (line 225) | def compute_throughout_metrics(batch: TensorDict, timing_raw: Dict[str, ... function _calculate_bootstrap_metrics (line 261) | def _calculate_bootstrap_metrics(group: pd.DataFrame, variable_name: str... function _process_prompt_group_task (line 329) | def _process_prompt_group_task(group: pd.DataFrame, numeric_variables: L... function bootstrap_metric (line 385) | def bootstrap_metric( function calc_maj_val (line 425) | def calc_maj_val(data: list[dict[str, Any]], vote_key: str, val_key: str... function process_validation_metrics (line 461) | def process_validation_metrics( function aggregate_validation_metrics (line 601) | def aggregate_validation_metrics(data_sources: List[str], sample_inputs:... FILE: siirl/utils/model_utils/activation_offload.py function _get_unique_tensor_key (line 32) | def _get_unique_tensor_key(tensor): class FSDPParameterFilter (line 37) | class FSDPParameterFilter: method __init__ (line 38) | def __init__(self): method __call__ (line 41) | def __call__(self, tensor): method update_model_parameters (line 44) | def update_model_parameters(self, model): class CpuOffloadHookWithOffloadHandler (line 51) | class CpuOffloadHookWithOffloadHandler: method __init__ (line 59) | def __init__( method __enter__ (line 70) | def __enter__(self): method __exit__ (line 74) | def __exit__(self, *args: Any): method on_save_for_backward (line 78) | def on_save_for_backward(self, tensor: torch.Tensor) -> Any: method on_get_saved_tensor (line 82) | def on_get_saved_tensor(self, saved_state: Any) -> torch.Tensor: class OffloadHandler (line 87) | class OffloadHandler: method __init__ (line 90) | def __init__(self) -> None: method tensor_push (line 93) | def tensor_push(self, tensor: torch.Tensor, **kwargs) -> Any: method tensor_pop (line 97) | def tensor_pop(self, tensor_tag: Any, **kwargs): class GroupCommitFunction (line 102) | class GroupCommitFunction(torch.autograd.Function): method forward (line 110) | def forward(ctx, tensor, cpu_offload_handler): method backward (line 118) | def backward(ctx, grad_output): class SynchronizedGroupOffloadHandler (line 128) | class SynchronizedGroupOffloadHandler(OffloadHandler): method __init__ (line 134) | def __init__(self, num_offload_group, tensor_need_offloading_checker=(... method groupid_reset (line 142) | def groupid_reset(self): method on_group_commit_forward (line 152) | def on_group_commit_forward(self): method on_group_commit_backward (line 158) | def on_group_commit_backward(self): method offload (line 164) | def offload(src_tensor, pin_memory=True): method reload (line 179) | def reload(state, non_blocking=None): method tensor_push (line 186) | def tensor_push(self, tensor: torch.Tensor, **kwargs): method tensor_pop (line 201) | def tensor_pop(self, tensor_tag, **kwargs): class AsyncDoubleBufferGroupOffloadHandler (line 212) | class AsyncDoubleBufferGroupOffloadHandler(SynchronizedGroupOffloadHandl... method __init__ (line 219) | def __init__( method tensor_push (line 254) | def tensor_push(self, tensor: torch.Tensor, **kwargs) -> Any: method tensor_pop (line 279) | def tensor_pop(self, tensor_tag, **kwargs): method bulk_offload_group (line 292) | def bulk_offload_group(self, group_to_offload): method synchronize_on_group_commit_forward (line 313) | def synchronize_on_group_commit_forward(self, current_group): method on_group_commit_forward (line 341) | def on_group_commit_forward(self): method bulk_reload_group (line 349) | def bulk_reload_group(self, group_to_reload): method on_group_commit_backward (line 367) | def on_group_commit_backward(self): function get_activation_offload_context (line 392) | def get_activation_offload_context(num_layers: int = 1, model_layers: in... class ActivationHandler (line 408) | class ActivationHandler: method __init__ (line 409) | def __init__(self, offload_ctx, sync_func, tensor_filter, enable_ckpt): method pre_forward (line 420) | def pre_forward(self, module): method post_forward (line 425) | def post_forward(self, module): method _pack_kwargs (line 429) | def _pack_kwargs(self, *args, **kwargs): method _unpack_kwargs (line 438) | def _unpack_kwargs(self, flat_args, kwarg_keys): method _ckpt_forward (line 446) | def _ckpt_forward(self, forward_method, *args, **kwargs): method forward (line 461) | def forward(self, module, forward_method, *args, **kwargs): method wrap_module_forward_method (line 477) | def wrap_module_forward_method(self, module): function enable_activation_offloading (line 492) | def enable_activation_offloading(model, strategy, enable_ckpt=False): FILE: siirl/utils/model_utils/attention_utils.py function _get_attention_functions (line 20) | def _get_attention_functions() -> tuple[Callable, Callable, Callable, Ca... function index_first_axis (line 37) | def index_first_axis(*args, **kwargs): function pad_input (line 53) | def pad_input(*args, **kwargs): function rearrange (line 69) | def rearrange(*args, **kwargs): function unpad_input (line 84) | def unpad_input(*args, **kwargs): FILE: siirl/utils/model_utils/flops_counter.py function get_device_flops (line 33) | def get_device_flops(unit="T"): class FlopsCounter (line 65) | class FlopsCounter: method __init__ (line 75) | def __init__(self, config: PretrainedConfig, forward_only: bool = False): method _estimate_unknown_flops (line 97) | def _estimate_unknown_flops(self, tokens_sum, batch_seqlens, delta_time): method _estimate_qwen2_flops (line 100) | def _estimate_qwen2_flops(self, tokens_sum, batch_seqlens, delta_time): method _estimate_internvl_flops (line 142) | def _estimate_internvl_flops(self, tokens_sum, batch_seqlens, delta_ti... method _estimate_deepseek_v3_flops (line 198) | def _estimate_deepseek_v3_flops(self, tokens_sum, batch_seqlens, delta... method _estimate_qwen3_moe_flops (line 258) | def _estimate_qwen3_moe_flops(self, tokens_sum, batch_seqlens, delta_t... method _estimate_openvla_flops (line 305) | def _estimate_openvla_flops(self, tokens_sum, batch_seqlens, delta_time): method estimate_flops (line 371) | def estimate_flops(self, batch_seqlens, delta_time): FILE: siirl/utils/model_utils/fsdp_utils.py function init_fn (line 45) | def init_fn(x: torch.nn.Module): function get_init_weight_context_manager (line 52) | def get_init_weight_context_manager(use_meta_tensor=True, mesh: DeviceMe... function get_fsdp_wrap_policy_vla (line 68) | def get_fsdp_wrap_policy_vla(module, config=None, is_lora=False): function get_fsdp_wrap_policy (line 185) | def get_fsdp_wrap_policy(module, config=None, is_lora=False): function offload_fsdp_model_to_cpu (line 268) | def offload_fsdp_model_to_cpu(model: FSDP, empty_cache: bool = True): function offload_fsdp2_model_to_cpu (line 291) | def offload_fsdp2_model_to_cpu(model, empty_cache: bool = True): function load_fsdp_model_to_gpu (line 299) | def load_fsdp_model_to_gpu(model: FSDP): function load_fsdp2_model_to_gpu (line 319) | def load_fsdp2_model_to_gpu(model): function offload_fsdp_optimizer (line 326) | def offload_fsdp_optimizer(optimizer): function load_fsdp_optimizer (line 338) | def load_fsdp_optimizer(optimizer, device_id): function meta_device_init (line 350) | def meta_device_init(): function parallel_load_safetensors (line 381) | def parallel_load_safetensors(filepath): function parallel_init_module_fn (line 437) | def parallel_init_module_fn(module: torch.nn.Module, shard_states: Dict[... function fsdp_version (line 516) | def fsdp_version(model): function get_fsdp_state_ctx (line 525) | def get_fsdp_state_ctx(model, state_type, state_cfg, optim_cfg): function fsdp2_load_full_state_dict (line 532) | def fsdp2_load_full_state_dict(model: torch.nn.Module, full_state: dict,... function apply_fsdp2 (line 563) | def apply_fsdp2(model, fsdp_kwargs, config): function fsdp2_clip_grad_norm_ (line 585) | def fsdp2_clip_grad_norm_(parameters, max_norm, norm_type=2.0, error_if_... function layered_summon_lora_params (line 601) | def layered_summon_lora_params(fsdp_module) -> OrderedDict: FILE: siirl/utils/model_utils/model.py class LambdaLayer (line 40) | class LambdaLayer(nn.Module): method __init__ (line 41) | def __init__(self, fn): method forward (line 45) | def forward(self, *args, **kwargs): function squeeze (line 49) | def squeeze(x): function update_model_config (line 53) | def update_model_config(module_config, override_config_kwargs): function get_huggingface_actor_config (line 66) | def get_huggingface_actor_config(model_name: str, override_config_kwargs... function get_generation_config (line 78) | def get_generation_config( function create_huggingface_actor (line 95) | def create_huggingface_actor(model_name: str, override_config_kwargs=Non... function create_huggingface_critic (line 119) | def create_huggingface_critic(model_name: str, override_config_kwargs=No... function get_model_size (line 141) | def get_model_size(model: nn.Module, scale="auto"): function print_model_size (line 168) | def print_model_size(model: nn.Module, name: str = None): function create_random_mask (line 175) | def create_random_mask( function compute_position_id_with_mask (line 218) | def compute_position_id_with_mask(mask): function normalize_model_name (line 222) | def normalize_model_name(name, pp_rank, vpp_rank, transformer_config, la... function normalize_pp_vpp_params (line 246) | def normalize_pp_vpp_params(params, num_hidden_layers, layer_name="layer... function get_parallel_model_from_config (line 267) | def get_parallel_model_from_config( function _get_parallel_model_architecture_from_config (line 285) | def _get_parallel_model_architecture_from_config(config: PretrainedConfi... function _load_hf_model (line 298) | def _load_hf_model(config, model_config, is_value_model, local_cache_path): function get_hf_model_path (line 352) | def get_hf_model_path(config, local_cache_path="~/.cache/siirl/rlhf"): function load_megatron_model_weights (line 365) | def load_megatron_model_weights( function load_megatron_gptmodel_weights (line 390) | def load_megatron_gptmodel_weights( function pad_packed_inputs (line 409) | def pad_packed_inputs(unpad_tokens: torch.Tensor, cu_seqlens, max_seqlen... function load_mcore_dist_weights (line 442) | def load_mcore_dist_weights(parallel_model, dist_weight_path, is_value_m... function get_parallel_gptmodel_from_config (line 463) | def get_parallel_gptmodel_from_config( function convert_weight_keys (line 499) | def convert_weight_keys(state_dict: Dict[str, torch.Tensor], model: PreT... function extract_multi_modal_inputs (line 520) | def extract_multi_modal_inputs( class CausalLMOutputForPPO (line 562) | class CausalLMOutputForPPO(CausalLMOutputWithPast): FILE: siirl/utils/model_utils/npu_utils.py class IndexFirstAxis (line 22) | class IndexFirstAxis(torch.autograd.Function): method forward (line 24) | def forward(ctx, input, indices): method backward (line 36) | def backward(ctx, grad_output): class IndexPutFirstAxis (line 56) | class IndexPutFirstAxis(torch.autograd.Function): method forward (line 58) | def forward(ctx, values, indices, first_axis_dim): method backward (line 69) | def backward(ctx, grad_output): function pad_input (line 81) | def pad_input(hidden_states, indices, batch, seqlen): function unpad_input (line 99) | def unpad_input(hidden_states, attention_mask, unused_mask=None): FILE: siirl/utils/model_utils/seqlen_balancing.py function calculate_workload (line 29) | def calculate_workload(seqlen_list: list[int]): function karmarkar_karp (line 41) | def karmarkar_karp(seqlen_list: list[int], k_partitions: int, equal_size... function greedy_partition (line 145) | def greedy_partition(seqlen_list: list[int], k_partitions: int, equal_si... function get_seqlen_balanced_partitions (line 165) | def get_seqlen_balanced_partitions(seqlen_list: list[int], k_partitions:... function log_seqlen_unbalance (line 209) | def log_seqlen_unbalance(seqlen_list: list[int], partitions: list[list[i... function ceildiv (line 257) | def ceildiv(a, b): function roundup_divisible (line 261) | def roundup_divisible(a, b): function rearrange_micro_batches (line 265) | def rearrange_micro_batches( function get_reverse_idx (line 341) | def get_reverse_idx(idx_map): function prepare_dynamic_batch (line 359) | def prepare_dynamic_batch( function restore_dynamic_batch (line 395) | def restore_dynamic_batch(data: torch.Tensor, batch_idx_list: list[list[... FILE: siirl/utils/model_utils/tensordict_utils.py function assign_non_tensor_dict (line 27) | def assign_non_tensor_dict(tensor_dict: TensorDict, non_tensor_dict: dict): function assign_non_tensor_data (line 33) | def assign_non_tensor_data(tensor_dict: TensorDict, key, val): function assign_non_tensor (line 37) | def assign_non_tensor(tensordict: TensorDict, **kwargs): function unwrap_non_tensor_data (line 43) | def unwrap_non_tensor_data(data): function get_non_tensor_data (line 49) | def get_non_tensor_data(data: TensorDict, key: str, default): function get_tensordict (line 54) | def get_tensordict(tensor_dict: dict[str, torch.Tensor | list], non_tens... function index_select_tensor_dict (line 99) | def index_select_tensor_dict(batch: TensorDict, indices: torch.Tensor | ... function union_tensor_dict (line 128) | def union_tensor_dict(tensor_dict1: TensorDict, tensor_dict2: TensorDict... function make_iterator (line 150) | def make_iterator(tensordict: TensorDict, mini_batch_size, epochs, seed=... function assert_tensordict_eq (line 176) | def assert_tensordict_eq(tensordict1: TensorDict, tensordict2: TensorDict): function pop (line 200) | def pop(tensordict: TensorDict, keys: Iterator[str]) -> TensorDict: function pad_to_divisor (line 216) | def pad_to_divisor(data: TensorDict, size_divisor: int): function unpad (line 244) | def unpad(data: TensorDict, pad_size): FILE: siirl/utils/model_utils/torch_dtypes.py class PrecisionType (line 27) | class PrecisionType: method supported_type (line 43) | def supported_type(precision: Union[str, int]) -> bool: method supported_types (line 47) | def supported_types() -> list[str]: method is_fp16 (line 51) | def is_fp16(precision): method is_fp32 (line 55) | def is_fp32(precision): method is_bf16 (line 59) | def is_bf16(precision): method to_dtype (line 63) | def to_dtype(precision): method to_str (line 74) | def to_str(precision): FILE: siirl/utils/model_utils/torch_functional.py function gather_from_labels (line 41) | def gather_from_labels(data, label): function logprobs_from_logits (line 56) | def logprobs_from_logits(logits, labels, inplace_backward=True): function logprobs_from_logits_flash_attn (line 85) | def logprobs_from_logits_flash_attn(logits, labels, inplace_backward=True): function logprobs_from_logits_naive (line 93) | def logprobs_from_logits_naive(logits, labels): function logprobs_from_logits_v2 (line 99) | def logprobs_from_logits_v2(logits: torch.FloatTensor, labels): function clip_by_value (line 119) | def clip_by_value(x, tensor_min, tensor_max): function entropy_from_logits (line 128) | def entropy_from_logits(logits: torch.Tensor): function entropy_from_logits_with_chunking (line 135) | def entropy_from_logits_with_chunking(logits: torch.Tensor, chunk_size: ... function masked_sum (line 146) | def masked_sum(values, mask, axis=None): function masked_mean (line 154) | def masked_mean(values, mask, axis=None): function masked_var (line 171) | def masked_var(values, mask, unbiased=True): function masked_whiten (line 189) | def masked_whiten(values, mask, shift_mean=True): function get_response_mask (line 209) | def get_response_mask(response_id: torch.Tensor, eos_token: Union[int, L... function get_eos_mask (line 232) | def get_eos_mask(response_id: torch.Tensor, eos_token: int = 2, dtype=to... function compute_grad_norm (line 257) | def compute_grad_norm(model: nn.Module): function broadcast_dict_tensor (line 265) | def broadcast_dict_tensor(tensors: Union[Dict[str, torch.Tensor], Tensor... function allgather_dict_tensors (line 275) | def allgather_dict_tensors(tensors: Union[Dict[str, torch.Tensor], Tenso... function split_dict_tensor_into_batches (line 309) | def split_dict_tensor_into_batches(tensors: TensorDict, batch_size) -> L... function pad_2d_list_to_length (line 316) | def pad_2d_list_to_length(response, pad_token_id, max_length=None): function pad_sequence_to_length (line 327) | def pad_sequence_to_length(tensors, max_seq_len, pad_token_id, left_pad=... function postprocess_data (line 340) | def postprocess_data( function tokenize_and_postprocess_data (line 393) | def tokenize_and_postprocess_data( function remove_pad_token (line 416) | def remove_pad_token(input_ids: torch.Tensor, attention_mask: torch.Tens... function log_probs_from_logits_response (line 431) | def log_probs_from_logits_response(input_ids, logits, response_length): function log_probs_from_logits_response_rmpad (line 447) | def log_probs_from_logits_response_rmpad(input_ids, attention_mask, logi... function log_probs_from_logits_all_rmpad (line 474) | def log_probs_from_logits_all_rmpad(input_ids_rmpad, logits_rmpad, indic... function post_process_logits (line 502) | def post_process_logits(input_ids, logits, temperature, top_k, top_p): function get_cosine_schedule_with_warmup (line 518) | def get_cosine_schedule_with_warmup( function get_constant_schedule_with_warmup (line 562) | def get_constant_schedule_with_warmup( function prepare_decoder_attention_mask (line 587) | def prepare_decoder_attention_mask(attention_mask, input_shape, inputs_e... function _make_causal_mask (line 611) | def _make_causal_mask(input_ids_shape: torch.Size, dtype: torch.dtype, d... function _expand_mask (line 624) | def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Option... function get_unpad_data (line 638) | def get_unpad_data(attention_mask): function get_wsd_schedule_with_warmup (line 650) | def get_wsd_schedule_with_warmup( function check_device_is_available (line 706) | def check_device_is_available(): function distributed_mean_max_min_std (line 718) | def distributed_mean_max_min_std(local_tensor, compute_max=True, compute... function distributed_masked_mean (line 761) | def distributed_masked_mean(local_tensor, local_mask): FILE: siirl/utils/model_utils/ulysses.py function set_ulysses_sequence_parallel_group (line 30) | def set_ulysses_sequence_parallel_group(group: dist.ProcessGroup): function get_ulysses_sequence_parallel_group (line 38) | def get_ulysses_sequence_parallel_group() -> Optional[dist.ProcessGroup]: function get_ulysses_sequence_parallel_world_size (line 46) | def get_ulysses_sequence_parallel_world_size(group: ProcessGroup = None)... function get_ulysses_sequence_parallel_rank (line 54) | def get_ulysses_sequence_parallel_rank(group: ProcessGroup = None) -> int: function gather_seq_scatter_heads (line 62) | def gather_seq_scatter_heads( function gather_heads_scatter_seq (line 86) | def gather_heads_scatter_seq(x: Tensor, head_dim: int, seq_dim: int, gro... function _pad_tensor (line 104) | def _pad_tensor(x: Tensor, dim: int, padding_size: int) -> Tensor: function _unpad_tensor (line 111) | def _unpad_tensor(x: Tensor, dim: int, padding_size: int) -> Tensor: function slice_input_tensor (line 117) | def slice_input_tensor(x: Tensor, dim: int, padding: bool = True, group:... function all_to_all_tensor (line 133) | def all_to_all_tensor( function all_gather_tensor (line 155) | def all_gather_tensor(local_tensor: Tensor, group: Optional[dist.Process... class SeqAllToAll (line 165) | class SeqAllToAll(torch.autograd.Function): method forward (line 167) | def forward( method backward (line 182) | def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor, No... class Gather (line 194) | class Gather(torch.autograd.Function): method forward (line 196) | def forward( method backward (line 224) | def backward(ctx: Any, grad_output: Tensor) -> Any: function gather_outpus_and_unpad (line 237) | def gather_outpus_and_unpad( function ulysses_pad (line 272) | def ulysses_pad(input_ids_rmpad: Tensor, position_ids_rmpad: Optional[Te... function ulysses_pad_and_slice_inputs (line 312) | def ulysses_pad_and_slice_inputs(input_ids_rmpad: torch.Tensor, position... function validate_ulysses_config (line 338) | def validate_ulysses_config(num_heads, ulysses_sequence_size): FILE: siirl/utils/model_utils/vllm_utils.py function patch_vllm_moe_model_weight_loader (line 64) | def patch_vllm_moe_model_weight_loader(model): class TensorLoRARequest (line 105) | class TensorLoRARequest(LoRARequest): class VLLMHijack (line 110) | class VLLMHijack: method hijack (line 112) | def hijack(): FILE: siirl/utils/reward_score/__init__.py function default_compute_score (line 19) | def default_compute_score(data_source, solution_str, ground_truth, extra... function _default_compute_score (line 102) | def _default_compute_score(data_source, solution_str, ground_truth, extr... FILE: siirl/utils/reward_score/embodied.py function _tensor_to_str_list (line 33) | def _tensor_to_str_list(tensor: torch.Tensor) -> List[str]: function _extract_task_name (line 41) | def _extract_task_name(task_file_name: str) -> str: function _compute_cluster_centers (line 47) | def _compute_cluster_centers(embeddings: np.ndarray, eps: float = 0.5, m... function _get_batch_size (line 71) | def _get_batch_size(batch_data: TensorDict) -> int: function _extract_local_data (line 90) | def _extract_local_data(batch_data: TensorDict) -> Dict[str, Any]: function _gather_all_data (line 123) | def _gather_all_data( function _compute_all_rewards (line 186) | def _compute_all_rewards(data: Dict[str, Any], logger) -> np.ndarray: function _build_results (line 259) | def _build_results( function compute_embodied_reward (line 287) | def compute_embodied_reward( FILE: siirl/utils/reward_score/geo3k.py function format_reward (line 23) | def format_reward(predict_str: str) -> float: function acc_reward (line 29) | def acc_reward(predict_str: str, ground_truth: str) -> float: function format_reward (line 34) | def format_reward(predict_str: str) -> float: function acc_reward (line 40) | def acc_reward(predict_str: str, ground_truth: str) -> float: function compute_score (line 45) | def compute_score(predict_str: str, ground_truth: str) -> float: FILE: siirl/utils/reward_score/gsm8k.py function extract_solution (line 18) | def extract_solution(solution_str, method="strict"): function compute_score (line 44) | def compute_score(solution_str, ground_truth, method="strict", format_sc... FILE: siirl/utils/reward_score/math.py function compute_score (line 17) | def compute_score(solution_str, ground_truth) -> float: function is_equiv (line 32) | def is_equiv(str1, str2, verbose=False): function remove_boxed (line 49) | def remove_boxed(s): function last_boxed_only_string (line 63) | def last_boxed_only_string(string): function fix_fracs (line 90) | def fix_fracs(string): function fix_a_slash_b (line 122) | def fix_a_slash_b(string): function remove_right_units (line 137) | def remove_right_units(string): function fix_sqrt (line 147) | def fix_sqrt(string): function strip_string (line 162) | def strip_string(string): FILE: siirl/utils/reward_score/math_batch.py function compute_score_batched (line 18) | def compute_score_batched(data_sources, solution_strs, ground_truths, ex... FILE: siirl/utils/reward_score/math_dapo.py function last_boxed_only_string (line 20) | def last_boxed_only_string(string: str) -> Optional[str]: function remove_boxed (line 50) | def remove_boxed(s: str) -> str: function normalize_final_answer (line 124) | def normalize_final_answer(final_answer: str) -> str: function is_correct_minerva (line 165) | def is_correct_minerva(solution_str: str, gt: str, gt_need_extract: bool... function is_correct_strict_box (line 191) | def is_correct_strict_box(pred: str, gt: str, pause_tokens_index: Option... function verify (line 216) | def verify(solution_str: str, answer: str, strict_box_verify: bool = Fal... function compute_score (line 236) | def compute_score( FILE: siirl/utils/reward_score/math_verify.py function compute_score (line 23) | def compute_score(model_output: str, ground_truth: str, timeout_score: f... FILE: siirl/utils/reward_score/mm_eureka.py function extract_answer_with_tags (line 13) | def extract_answer_with_tags(text): function accuracy_reward_func (line 20) | def accuracy_reward_func(completion, answer): function format_reward_func (line 64) | def format_reward_func(completion, **kwargs): function compute_score (line 78) | def compute_score(predict_str: str, ground_truth: str) -> float: FILE: siirl/utils/reward_score/prime_code/__init__.py function compute_score (line 21) | def compute_score(completion, test_cases, continuous=False): FILE: siirl/utils/reward_score/prime_code/testing_util.py function truncatefn (line 39) | def truncatefn(s, length=300): class CODE_TYPE (line 47) | class CODE_TYPE(Enum): class Capturing (line 55) | class Capturing(list): method __enter__ (line 56) | def __enter__(self): method __exit__ (line 63) | def __exit__(self, *args): function only_int_check (line 69) | def only_int_check(val): function string_int_check (line 73) | def string_int_check(val): function combined_int_check (line 77) | def combined_int_check(val): function clean_traceback (line 81) | def clean_traceback(error_traceback): function run_test (line 88) | def run_test(in_outs, test=None, debug=False, timeout=15): function custom_compare_ (line 510) | def custom_compare_(output, ground_truth): function stripped_string_compare (line 525) | def stripped_string_compare(s1, s2): function call_method (line 531) | def call_method(method, inputs): function reliability_guard (line 557) | def reliability_guard(maximum_memory_bytes=None): FILE: siirl/utils/reward_score/prime_code/utils.py function _temp_run (line 26) | def _temp_run(sample, generation, debug, result, metadata_list, timeout): function check_correctness (line 41) | def check_correctness(in_outs: Optional[dict], generation, timeout=10, d... FILE: siirl/utils/reward_score/prime_math/__init__.py function _sympy_parse (line 44) | def _sympy_parse(expr: str): function _parse_latex (line 53) | def _parse_latex(expr: str) -> str: function _is_float (line 71) | def _is_float(num: str) -> bool: function _is_int (line 79) | def _is_int(x: float) -> bool: function _is_frac (line 86) | def _is_frac(expr: str) -> bool: function _str_is_int (line 90) | def _str_is_int(x: str) -> bool: function _str_to_int (line 99) | def _str_to_int(x: str) -> bool: function _inject_implicit_mixed_number (line 105) | def _inject_implicit_mixed_number(step: str): function _strip_properly_formatted_commas (line 115) | def _strip_properly_formatted_commas(expr: str): function _normalize (line 126) | def _normalize(expr: str) -> str: function count_unknown_letters_in_expr (line 193) | def count_unknown_letters_in_expr(expr: str): function should_allow_eval (line 200) | def should_allow_eval(expr: str): function are_equal_under_sympy (line 213) | def are_equal_under_sympy(ground_truth_normalized: str, given_normalized... function split_tuple (line 227) | def split_tuple(expr: str): function grade_answer (line 241) | def grade_answer(given_answer: str, ground_truth: str) -> bool: function remove_boxed (line 297) | def remove_boxed(s): function _last_boxed_only_string (line 307) | def _last_boxed_only_string(string): function match_answer (line 337) | def match_answer(response): function compute_score (line 379) | def compute_score(model_output: str, ground_truth: str) -> bool: FILE: siirl/utils/reward_score/prime_math/grader.py function is_digit (line 110) | def is_digit(s): function normalize (line 122) | def normalize(answer, pi) -> str: function handle_base (line 140) | def handle_base(x) -> str: function handle_pi (line 149) | def handle_pi(string, pi): function math_equal (line 173) | def math_equal( function symbolic_equal (line 282) | def symbolic_equal(a, b, tolerance, timeout=10.0): function format_intervals (line 320) | def format_intervals(prediction): FILE: siirl/utils/reward_score/prime_math/math_normalize.py function normalize_answer (line 44) | def normalize_answer(answer: Optional[str]) -> Optional[str]: function _fix_fracs (line 58) | def _fix_fracs(string): function _fix_a_slash_b (line 90) | def _fix_a_slash_b(string): function _remove_right_units (line 105) | def _remove_right_units(string): function _fix_sqrt (line 115) | def _fix_sqrt(string): function _strip_string (line 130) | def _strip_string(string): FILE: siirl/utils/reward_score/sandbox_fusion/__init__.py function compute_score (line 28) | def compute_score(sandbox_fusion_url, concurrent_semaphore, completion, ... FILE: siirl/utils/reward_score/sandbox_fusion/utils.py function call_sandbox_api (line 37) | def call_sandbox_api(sandbox_fusion_url: str, code: str, stdin: str, com... function _process_single_case (line 128) | def _process_single_case(case_index: int, stdin_data: Any, expected_outp... function check_correctness (line 367) | def check_correctness(sandbox_fusion_url: str, in_outs: Optional[dict], ... FILE: siirl/utils/reward_score/search_r1_like_qa_em.py function normalize_answer (line 23) | def normalize_answer(s): function em_check (line 40) | def em_check(prediction, golden_answers): function subem_check (line 53) | def subem_check(prediction, golden_answers): function extract_solution (line 66) | def extract_solution(solution_str): function count_answer_tags (line 89) | def count_answer_tags(text): function compute_score (line 96) | def compute_score(solution_str, ground_truth, method="strict", format_sc... function compute_score_subem (line 131) | def compute_score_subem(solution_str, ground_truth, method="strict", for... FILE: tests/dag/test_config_loader.py function capture_loguru_logs (line 31) | def capture_loguru_logs(level="INFO"): class TestDAGConfigLoader (line 66) | class TestDAGConfigLoader(unittest.TestCase): method setUpClass (line 72) | def setUpClass(cls): method tearDownClass (line 85) | def tearDownClass(cls): method setUp (line 98) | def setUp(self): method tearDown (line 106) | def tearDown(self): method _create_config_file (line 115) | def _create_config_file(self, content: str, file_type: str = "yaml"): method test_load_valid_yaml_config (line 122) | def test_load_valid_yaml_config(self): method test_load_valid_json_config (line 181) | def test_load_valid_json_config(self): method test_missing_dag_id (line 200) | def test_missing_dag_id(self): method test_missing_nodes_list (line 213) | def test_missing_nodes_list(self): method test_empty_nodes_list (line 223) | def test_empty_nodes_list(self): method test_missing_node_id (line 235) | def test_missing_node_id(self): method test_missing_node_type (line 247) | def test_missing_node_type(self): method test_invalid_node_type (line 259) | def test_invalid_node_type(self): method test_invalid_node_role (line 272) | def test_invalid_node_role(self): method test_duplicate_node_ids (line 286) | def test_duplicate_node_ids(self): method test_non_existent_dependency (line 302) | def test_non_existent_dependency(self): method test_circular_dependency (line 315) | def test_circular_dependency(self): method test_executable_ref_not_found_module (line 334) | def test_executable_ref_not_found_module(self): method test_executable_ref_not_found_function (line 348) | def test_executable_ref_not_found_function(self): method test_executable_ref_not_callable (line 362) | def test_executable_ref_not_callable(self): method test_ref_resolution (line 376) | def test_ref_resolution(self): method test_ref_resolution_invalid_path (line 403) | def test_ref_resolution_invalid_path(self): method test_ref_resolution_invalid_path_nested (line 421) | def test_ref_resolution_invalid_path_nested(self): method test_node_role_validation (line 440) | def test_node_role_validation(self): method test_save_dag_pic_no_nodes (line 482) | def test_save_dag_pic_no_nodes(self, mock_digraph): method test_save_dag_pic_invalid_graph (line 497) | def test_save_dag_pic_invalid_graph(self, mock_digraph): method test_load_profiler_yaml_config (line 529) | def test_load_profiler_yaml_config(self): method test_profiler_npu_environment (line 574) | def test_profiler_npu_environment(self): FILE: tests/dag/test_node.py function local_sync_executable (line 23) | def local_sync_executable(data: Any, node_config: Dict = None) -> Dict: function local_async_executable (line 29) | async def local_async_executable(data: Any, node_config: Dict = None) ->... function always_fail_func (line 41) | async def always_fail_func(**kwargs): class TestNode (line 47) | class TestNode(unittest.TestCase): method test_node_creation_valid (line 50) | def test_node_creation_valid(self): method test_node_creation_minimal (line 62) | def test_node_creation_minimal(self): method test_node_creation_invalid_id (line 73) | def test_node_creation_invalid_id(self): method test_node_creation_invalid_type (line 78) | def test_node_creation_invalid_type(self): method test_resolve_executable_non_existent (line 83) | def test_resolve_executable_non_existent(self): method test_resolve_executable_not_callable (line 88) | def test_resolve_executable_not_callable(self): method test_add_remove_dependency (line 93) | def test_add_remove_dependency(self): method test_is_ready (line 107) | def test_is_ready(self): method test_update_status (line 123) | def test_update_status(self): method test_retry_logic (line 139) | def test_retry_logic(self): method test_execute_no_executable (line 152) | def test_execute_no_executable(self): method test_execute_sync_executable (line 164) | def test_execute_sync_executable(self): method test_execute_async_executable (line 180) | def test_execute_async_executable(self): method test_execute_failing_executable (line 193) | def test_execute_failing_executable(self): FILE: tests/dag/test_task_graph.py function example_data_load_func (line 20) | def example_data_load_func(): function example_compute_func (line 24) | def example_compute_func(): class TestTaskGraph (line 28) | class TestTaskGraph(unittest.TestCase): method setUp (line 31) | def setUp(self): method test_graph_creation (line 42) | def test_graph_creation(self): method test_add_node (line 46) | def test_add_node(self): method test_build_adjacency_lists (line 56) | def test_build_adjacency_lists(self): method test_get_node (line 78) | def test_get_node(self): method test_get_dependencies_and_dependents (line 83) | def test_get_dependencies_and_dependents(self): method test_get_entry_and_exit_nodes (line 95) | def test_get_entry_and_exit_nodes(self): method test_validate_graph_valid (line 107) | def test_validate_graph_valid(self): method test_validate_graph_missing_dependency (line 115) | def test_validate_graph_missing_dependency(self): method test_validate_graph_cyclic (line 123) | def test_validate_graph_cyclic(self): method test_get_topological_sort_valid (line 133) | def test_get_topological_sort_valid(self): method test_get_topological_sort_empty_graph (line 149) | def test_get_topological_sort_empty_graph(self): method test_get_topological_sort_cyclic (line 152) | def test_get_topological_sort_cyclic(self): method test_reset_nodes_status (line 161) | def test_reset_nodes_status(self): method test_load_from_config_valid (line 178) | def test_load_from_config_valid(self): method test_load_from_config_missing_field (line 195) | def test_load_from_config_missing_field(self): method test_load_from_config_invalid_enum_value (line 203) | def test_load_from_config_invalid_enum_value(self): method test_load_from_config_invalid_graph_structure (line 211) | def test_load_from_config_invalid_graph_structure(self): FILE: tests/dag/test_task_loader.py function get_signatures (line 25) | def get_signatures(graphs: List[TaskGraph]) -> List[str]: class TestTaskLoaderInternals (line 30) | class TestTaskLoaderInternals(unittest.TestCase): method setUp (line 33) | def setUp(self): method test_generate_structural_signature (line 59) | def test_generate_structural_signature(self): method test_get_all_downstream_nodes_recursive (line 78) | def test_get_all_downstream_nodes_recursive(self): method test_get_all_ancestors (line 104) | def test_get_all_ancestors(self): method test_find_all_paths (line 126) | def test_find_all_paths(self): class TestTaskSplittingLogic (line 157) | class TestTaskSplittingLogic(unittest.TestCase): method assertGraphStructure (line 160) | def assertGraphStructure(self, graph: TaskGraph, expected_node_ids: Se... method test_split_single_structure_reconvergence (line 167) | def test_split_single_structure_reconvergence(self): method test_split_by_fan_out_to_exits (line 212) | def test_split_by_fan_out_to_exits(self): method test_split_by_reconverging_paths_diamond (line 245) | def test_split_by_reconverging_paths_diamond(self): function add_nodes_and_build (line 271) | def add_nodes_and_build(self, nodes: List[Node]) -> TaskGraph: class TestDiscoverAndSplitParallelPaths (line 280) | class TestDiscoverAndSplitParallelPaths(unittest.TestCase): method assertListOfGraphStructuresEqual (line 283) | def assertListOfGraphStructuresEqual(self, actual_graphs: List[TaskGra... method test_empty_graph (line 299) | def test_empty_graph(self): method test_linear_graph (line 304) | def test_linear_graph(self): method test_simple_reconvergence_diamond_graph (line 317) | def test_simple_reconvergence_diamond_graph(self): method test_fan_out_only_graph (line 334) | def test_fan_out_only_graph(self): method test_ex1_reconverge_from_prompt (line 347) | def test_ex1_reconverge_from_prompt(self): method test_ex2_complex_from_prompt (line 370) | def test_ex2_complex_from_prompt(self): FILE: tests/dag_worker/test_dag_worker.py class MockTokenizer (line 28) | class MockTokenizer: method __init__ (line 29) | def __init__(self, pad_token_id=0): class MockDAGWorker (line 35) | class MockDAGWorker: method __init__ (line 41) | def __init__(self, rank: int, world_size: int, data_buffers: List[ray.... method put_data_to_buffers (line 55) | async def put_data_to_buffers(self, key: str, data: TensorDict, source... method get_data_from_buffers (line 75) | async def get_data_from_buffers(self, key: str, my_current_dp_rank: in... method get_rank (line 102) | def get_rank(self): method barrier (line 105) | def barrier(self): function _create_test_dp (line 109) | def _create_test_dp(batch_size: int, seq_len: int, meta: Optional[Dict] ... class TestDAGWorkerDataFlow (line 119) | class TestDAGWorkerDataFlow(unittest.IsolatedAsyncioTestCase): method setUpClass (line 121) | def setUpClass(cls): method tearDownClass (line 126) | def tearDownClass(cls): method test_put_get_flow_sharded (line 130) | async def test_put_get_flow_sharded(self): method test_put_get_flow_object_ref (line 187) | async def test_put_get_flow_object_ref(self): FILE: tests/dag_worker/test_dapo_merge.py function get_unpacked_data (line 12) | def get_unpacked_data(value): function create_mock_batch (line 19) | def create_mock_batch(batch_size=1024): function test_filter_tensordict (line 90) | def test_filter_tensordict(): function test_merge_cached_and_new (line 123) | def test_merge_cached_and_new(): function test_full_dapo_workflow (line 262) | def test_full_dapo_workflow(): FILE: tests/dag_worker/test_dapo_pipeline.py function test_preprocess_dataloader (line 14) | def test_preprocess_dataloader(): class MockDAGWorker (line 53) | class MockDAGWorker: method __init__ (line 54) | def __init__(self, rank): method postprocess_sampling (line 59) | def postprocess_sampling(self, config, batch: TensorDict, filtered_ind... function sample_batch (line 136) | def sample_batch(): function test_postprocess_sampling_caching (line 145) | def test_postprocess_sampling_caching(sample_batch): function test_postprocess_sampling_merging (line 161) | def test_postprocess_sampling_merging(sample_batch): FILE: tests/data_buffer/detailed_put_performance_test.py function log_with_time (line 34) | def log_with_time(message: str): function create_mock_sample (line 40) | def create_mock_sample() -> TensorDict: function producer_task_detailed_profile (line 50) | async def producer_task_detailed_profile( function analyze_timings (line 82) | def analyze_timings(operation_name: str, timings: List[float]): function main (line 98) | async def main(): FILE: tests/data_buffer/performance_test_data_buffer.py function log_with_time (line 52) | def log_with_time(message: str): function create_mock_sample (line 58) | def create_mock_sample(item_idx: int) -> TensorDict: function producer_task (line 68) | async def producer_task( function main (line 94) | async def main(): FILE: tests/data_buffer/test_data_buffer.py class TestDataCoordinator (line 28) | class TestDataCoordinator(unittest.IsolatedAsyncioTestCase): method setUpClass (line 34) | def setUpClass(cls): method tearDownClass (line 39) | def tearDownClass(cls): method asyncSetUp (line 43) | async def asyncSetUp(self): method asyncTearDown (line 50) | async def asyncTearDown(self): method _create_mock_sample (line 57) | def _create_mock_sample(self, content_id: int) -> TensorDict: method _create_mock_sample_info (line 61) | def _create_mock_sample_info(self, tokens: int, group: int = 0) -> Sam... method test_put_increases_size (line 73) | async def test_put_increases_size(self): method test_get_batch_simple (line 86) | async def test_get_batch_simple(self): method test_get_batch_insufficient_samples (line 109) | async def test_get_batch_insufficient_samples(self): method test_get_batch_fifo_order (line 125) | async def test_get_batch_fifo_order(self): method test_get_batch_with_filter (line 143) | async def test_get_batch_with_filter(self): method test_get_batch_with_node_affinity_filter (line 179) | async def test_get_batch_with_node_affinity_filter(self): FILE: tests/scheduler/test_process_group_manager.py class TestProcessGroupManager (line 22) | class TestProcessGroupManager(unittest.TestCase): method _create_task_graph (line 23) | def _create_task_graph(self, graph_id: str, node_specs: List[tuple[str... method test_initialization_default_relevant_types (line 44) | def test_initialization_default_relevant_types(self): method test_initialization_custom_relevant_types (line 48) | def test_initialization_custom_relevant_types(self): method test_initialization_invalid_custom_relevant_types (line 53) | def test_initialization_invalid_custom_relevant_types(self): method test_default_filtering_model_nodes_only (line 59) | def test_default_filtering_model_nodes_only(self): method test_custom_filtering_compute_and_load_only (line 100) | def test_custom_filtering_compute_and_load_only(self): method test_no_relevant_nodes_in_graph_custom_filter (line 136) | def test_no_relevant_nodes_in_graph_custom_filter(self): method test_empty_relevant_types_set_no_nodes_processed (line 155) | def test_empty_relevant_types_set_no_nodes_processed(self): method test_complex_scenario_with_dependencies_and_default_filtering (line 173) | def test_complex_scenario_with_dependencies_and_default_filtering(self): method test_highly_complex_scenario_shared_nodes_16_workers (line 276) | def test_highly_complex_scenario_shared_nodes_16_workers(self): FILE: tests/scheduler/test_task_scheduler.py function create_test_graph (line 25) | def create_test_graph(graph_id: str, num_nodes: int, model_params: any =... class TestTaskScheduler (line 70) | class TestTaskScheduler(unittest.TestCase): method setUp (line 75) | def setUp(self): method test_scheduler_initialization (line 86) | def test_scheduler_initialization(self): method test_reset_scheduler_state (line 103) | def test_reset_scheduler_state(self): method test_get_original_graph_id (line 130) | def test_get_original_graph_id(self): method test_apportion_workers_to_tasks_even_strategy (line 145) | def test_apportion_workers_to_tasks_even_strategy(self): method test_apportion_workers_to_tasks_param_aware_strategy (line 187) | def test_apportion_workers_to_tasks_param_aware_strategy(self): method test_schedule_no_tasks (line 229) | def test_schedule_no_tasks(self): method test_schedule_fewer_tasks_than_workers (line 237) | def test_schedule_fewer_tasks_than_workers(self): method test_schedule_more_tasks_than_workers_raises_error (line 272) | def test_schedule_more_tasks_than_workers_raises_error(self): method test_schedule_with_task_splitting (line 283) | def test_schedule_with_task_splitting(self): method test_placement_logic_node_cohesion (line 331) | def test_placement_logic_node_cohesion(self): method test_placement_no_node_load_no_rank_preference (line 392) | def test_placement_no_node_load_no_rank_preference(self):