trinity.common.models package

Submodules

• trinity.common.models.mm_utils module
  • build_multi_modal_inputs()
  • attach_images_to_messages()
• trinity.common.models.model module
  • InferenceModel
    • InferenceModel.generate()
    • InferenceModel.chat()
    • InferenceModel.logprobs()
    • InferenceModel.convert_messages_to_experience()
    • InferenceModel.get_model_version()
    • InferenceModel.get_available_address()
  • ModelWrapper
    • ModelWrapper.__init__()
    • ModelWrapper.generate()
    • ModelWrapper.generate_async()
    • ModelWrapper.generate_mm()
    • ModelWrapper.generate_mm_async()
    • ModelWrapper.chat()
    • ModelWrapper.chat_async()
    • ModelWrapper.chat_mm()
    • ModelWrapper.chat_mm_async()
    • ModelWrapper.logprobs()
    • ModelWrapper.logprobs_async()
    • ModelWrapper.convert_messages_to_experience()
    • ModelWrapper.convert_messages_to_experience_async()
    • ModelWrapper.model_version
    • ModelWrapper.get_openai_client()
    • ModelWrapper.get_openai_async_client()
    • ModelWrapper.extract_experience_from_history()
  • convert_api_output_to_experience()
  • extract_logprobs()
• trinity.common.models.utils module
  • tokenize_and_mask_messages_hf()
  • tokenize_and_mask_messages_default()
  • get_action_mask_method()
  • get_checkpoint_dir_with_step_num()
  • load_state_dict()
  • merge_by_placement()
  • get_verl_checkpoint_info()
  • load_fsdp_state_dict_from_verl_checkpoint()
  • load_huggingface_state_dict()
  • get_megatron_converter()
• trinity.common.models.vllm_model module
  • vLLMRolloutModel
    • vLLMRolloutModel.__init__()
    • vLLMRolloutModel.chat()
    • vLLMRolloutModel.generate()
    • vLLMRolloutModel.chat_mm()
    • vLLMRolloutModel.generate_mm()
    • vLLMRolloutModel.logprobs()
    • vLLMRolloutModel.convert_messages_to_experience()
    • vLLMRolloutModel.shutdown()
    • vLLMRolloutModel.sync_model()
    • vLLMRolloutModel.init_process_group()
    • vLLMRolloutModel.run_api_server()
    • vLLMRolloutModel.has_api_server()
    • vLLMRolloutModel.api_server_ready()
    • vLLMRolloutModel.reset_prefix_cache()
    • vLLMRolloutModel.get_model_version()
    • vLLMRolloutModel.sleep()
    • vLLMRolloutModel.wake_up()
• trinity.common.models.vllm_worker module
  • WorkerExtension
    • WorkerExtension.init_process_group()
    • WorkerExtension.update_weight()

Module contents

trinity.common.models.create_inference_models(config: Config) → Tuple[List[InferenceModel], List[List[InferenceModel]]]

Create engine_num rollout models.

Each model has tensor_parallel_size workers.