from verl.utils.torch_functional import get_response_mask, pad_sequence_to_length
from verl.utils.device import get_visible_devices_keyword, get_device_name
super().__init__()
self.config = config
self.attention_backend = self.config.engine_kwargs.sglang.attention_backend
# get tp_rank of this process in this tp group
visible_devices = [None] * self._device_mesh_cpu.size(1)
devices_keyword = get_visible_devices_keyword()
torch.distributed.all_gather_object(
visible_devices, os.environ[devices_keyword], self._device_mesh_cpu.get_group("tp")
)
self.visible_devices_set = set(",".join(visible_devices).split(","))
os.environ[devices_keyword] = ",".join(sorted(list(self.visible_devices_set)))
node_rank = self._tp_rank // tp_size_per_node
first_rank_in_node = self._tp_rank % tp_size_per_node == 0
enable_memory_saver = True
if get_device_name() == "npu":
enable_memory_saver = False
backend = "fa3" if self.attention_backend != "ascend" else "ascend"
if first_rank_in_node:
rank = dist.get_rank()
os.environ["SGLANG_BLOCK_NONZERO_RANK_CHILDREN"] = "0"
self._engine = AsyncEngine(
model_path=actor_module,
dtype=self.config.dtype,
mem_fraction_static=self.config.gpu_memory_utilization,
enable_memory_saver=enable_memory_saver,
base_gpu_id=0,
gpu_id_step=1,
tp_size=self._tp_size,
node_rank=node_rank,
load_format=load_format,
dist_init_addr=dist_init_addr,
nnodes=nnodes,
trust_remote_code=trust_remote_code,
# NOTE(linjunrong): add rank to prevent SGLang generate same port inside PortArgs.init_new
# when random.seed is being set during training
port=30000 + rank,
# NOTE(Chenyang): if you want to debug the SGLang engine output
# please set the following parameters
# Otherwise, it will make the engine run too slow
# log_level="INFO",
# log_requests=True,
# log_requests_level=2,
# max_running_requests=1,
mm_attention_backend=backend,
attention_backend=backend,
# In async mode, we want token in token out.
skip_tokenizer_init=self.config.mode == "async",
)
