'
+ img_str_list.append(img_str)
+ elif self.is_video_file(img):
+ ret.append(((img,), None))
+
+ msg = msg.strip()
+ img_place_holder = ""
+ for img_str in img_str_list:
+ img_place_holder += f"{img_str}\n\n"
+
+ if len(img_str_list) > 0:
+ msg = f"{img_place_holder}\n\n{msg}"
+
+ if len(msg) > 0:
+ ret.append([msg, None])
+ else:
+ ret.append([msg, None])
+ else:
+ ret[-1][-1] = msg
+ return ret
+
+ def copy(self):
+ return Conversation(system=self.system, roles=self.roles, messages=[[x, y] for x, y in self.messages], offset=self.offset, sep_style=self.sep_style, sep=self.sep, sep2=self.sep2, version=self.version)
+
+ def dict(self):
+ if len(self.get_images()) > 0:
+ return {
+ "system": self.system,
+ "roles": self.roles,
+ "messages": [[x, y[0] if type(y) is tuple else y] for x, y in self.messages],
+ "offset": self.offset,
+ "sep": self.sep,
+ "sep2": self.sep2,
+ }
+ return {
+ "system": self.system,
+ "roles": self.roles,
+ "messages": self.messages,
+ "offset": self.offset,
+ "sep": self.sep,
+ "sep2": self.sep2,
+ }
+
+
+conv_vicuna_v0 = Conversation(
+ system="A chat between a curious human and an artificial intelligence assistant. " "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+ roles=("Human", "Assistant"),
+ messages=[
+ ["Human", "What are the key differences between renewable and non-renewable energy sources?"],
+ [
+ "Assistant",
+ "Renewable energy sources are those that can be replenished naturally in a relatively "
+ "short amount of time, such as solar, wind, hydro, geothermal, and biomass. "
+ "Non-renewable energy sources, on the other hand, are finite and will eventually be "
+ "depleted, such as coal, oil, and natural gas. Here are some key differences between "
+ "renewable and non-renewable energy sources:\n"
+ "1. Availability: Renewable energy sources are virtually inexhaustible, while non-renewable "
+ "energy sources are finite and will eventually run out.\n"
+ "2. Environmental impact: Renewable energy sources have a much lower environmental impact "
+ "than non-renewable sources, which can lead to air and water pollution, greenhouse gas emissions, "
+ "and other negative effects.\n"
+ "3. Cost: Renewable energy sources can be more expensive to initially set up, but they typically "
+ "have lower operational costs than non-renewable sources.\n"
+ "4. Reliability: Renewable energy sources are often more reliable and can be used in more remote "
+ "locations than non-renewable sources.\n"
+ "5. Flexibility: Renewable energy sources are often more flexible and can be adapted to different "
+ "situations and needs, while non-renewable sources are more rigid and inflexible.\n"
+ "6. Sustainability: Renewable energy sources are more sustainable over the long term, while "
+ "non-renewable sources are not, and their depletion can lead to economic and social instability.\n",
+ ],
+ ],
+ offset=2,
+ sep_style=SeparatorStyle.SINGLE,
+ sep="###",
+)
+
+conv_vicuna_v1 = Conversation(
+ system="A chat between a curious user and an artificial intelligence assistant. " "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+ roles=("USER", "ASSISTANT"),
+ version="v1",
+ messages=[],
+ offset=0,
+ sep_style=SeparatorStyle.TWO,
+ sep=" ",
+ sep2="",
+)
+
+conv_llama_2 = Conversation(
+ system="""You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
+
+If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.""",
+ roles=("USER", "ASSISTANT"),
+ version="llama_v2",
+ messages=[],
+ offset=0,
+ sep_style=SeparatorStyle.LLAMA_2,
+ sep="Unmask Viewer
+ +
+
+
+
+
+
+
+"""
+
+out_path = "trace_viewer.html"
+with open(out_path, "w", encoding="utf-8") as f:
+ f.write(html)
+
+print(f"[Trace Viewer] Wrote: {os.path.abspath(out_path)} (open in your browser)")
+
+
+
diff --git a/sample/jetengine_ext/__init__.py b/sample/jetengine_ext/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..4725d2d06f87d0ffa8c6a6d9344b2180ddbc7836
--- /dev/null
+++ b/sample/jetengine_ext/__init__.py
@@ -0,0 +1,8 @@
+from jetengine_ext.llm import LLM
+from jetengine_ext.sampling_params import SamplingParams
+
+
+__all__ = [
+ "LLM",
+ "SamplingParams",
+]
diff --git a/sample/jetengine_ext/config.py b/sample/jetengine_ext/config.py
new file mode 100644
index 0000000000000000000000000000000000000000..916b46773539ef2ffe45e6b3640e99fedb9d5fe2
--- /dev/null
+++ b/sample/jetengine_ext/config.py
@@ -0,0 +1,29 @@
+import os
+from dataclasses import dataclass
+from transformers import AutoConfig
+
+
+@dataclass
+class Config:
+ model: str
+ max_num_batched_tokens: int = 16384
+ max_num_seqs: int = 512
+ max_model_len: int = 4096
+ gpu_memory_utilization: float = 0.5
+ tensor_parallel_size: int = 1
+ enforce_eager: bool = False
+ hf_config: AutoConfig | None = None
+ eos: int = -1
+ kvcache_block_size: int = 256
+ num_kvcache_blocks: int = -1
+ mask_token_id: int = -1
+ block_length: int = 4
+
+ def __post_init__(self):
+ assert os.path.isdir(self.model)
+ assert self.kvcache_block_size % 256 == 0
+ assert 1 <= self.tensor_parallel_size <= 8
+ self.hf_config = AutoConfig.from_pretrained(self.model, trust_remote_code=True)
+ self.max_model_len = min(self.max_model_len, self.hf_config.max_position_embeddings)
+ assert self.max_num_batched_tokens >= self.max_model_len
+ assert self.mask_token_id != -1, "Mask token ID must be set"
diff --git a/sample/jetengine_ext/layers/activation.py b/sample/jetengine_ext/layers/activation.py
new file mode 100644
index 0000000000000000000000000000000000000000..83e4734a80570c5daef733dfffdc3dcca8abc3dc
--- /dev/null
+++ b/sample/jetengine_ext/layers/activation.py
@@ -0,0 +1,15 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+from liger_kernel.ops.swiglu import LigerSiLUMulFunction
+
+
+class SiluAndMul(nn.Module):
+
+ def __init__(self):
+ super().__init__()
+
+ @torch.compile
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x, y = x.chunk(2, -1)
+ return LigerSiLUMulFunction.apply(x, y)
diff --git a/sample/jetengine_ext/layers/attention.py b/sample/jetengine_ext/layers/attention.py
new file mode 100644
index 0000000000000000000000000000000000000000..f9ffae6c906e7cbb1251bc13041c4032062d41c4
--- /dev/null
+++ b/sample/jetengine_ext/layers/attention.py
@@ -0,0 +1,103 @@
+import torch
+from torch import nn
+import triton
+import triton.language as tl
+
+from jetengine_ext.utils.context import get_context
+from jetengine_ext.engine.sequence import RunType
+from jetengine_ext.kernels.triton.attention import sparse_attn_varlen
+# from jetengine.kernels.triton.attention import fused_kv_cache_attention
+# from jetengine.kernels.triton.attention import fused_kv_cache_attention_v5
+from flash_attn import flash_attn_with_kvcache
+
+
+@triton.jit
+def store_kvcache_kernel(
+ key_ptr,
+ key_stride,
+ value_ptr,
+ value_stride,
+ k_cache_ptr,
+ v_cache_ptr,
+ slot_mapping_ptr,
+ D: tl.constexpr,
+):
+ idx = tl.program_id(0)
+ key_offsets = idx * key_stride + tl.arange(0, D)
+ value_offsets = idx * value_stride + tl.arange(0, D)
+ key = tl.load(key_ptr + key_offsets)
+ value = tl.load(value_ptr + value_offsets)
+ slot = tl.load(slot_mapping_ptr + idx)
+ cache_offsets = slot * D + tl.arange(0, D)
+ tl.store(k_cache_ptr + cache_offsets, key)
+ tl.store(v_cache_ptr + cache_offsets, value)
+
+
+def store_kvcache(key: torch.Tensor, value: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, slot_mapping: torch.Tensor):
+ N, num_heads, head_dim = key.shape
+ D = num_heads * head_dim
+ assert key.stride(-1) == 1 and value.stride(-1) == 1
+ assert key.stride(1) == head_dim and value.stride(1) == head_dim
+ assert k_cache.stride(1) == D and v_cache.stride(1) == D
+ assert slot_mapping.numel() == N
+ store_kvcache_kernel[(N,)](key, key.stride(0), value, value.stride(0), k_cache, v_cache, slot_mapping, D)
+
+
+class Attention(nn.Module):
+
+ def __init__(
+ self,
+ num_heads,
+ head_dim,
+ scale,
+ num_kv_heads,
+ ):
+ super().__init__()
+ self.num_heads = num_heads
+ self.head_dim = head_dim
+ self.scale = scale
+ self.num_kv_heads = num_kv_heads
+ self.k_cache = self.v_cache = torch.tensor([])
+
+ def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
+ pass
+
+class BlockAttention(Attention):
+ def __init__(
+ self,
+ num_heads,
+ head_dim,
+ scale,
+ num_kv_heads,
+ ):
+ super().__init__(num_heads, head_dim, scale, num_kv_heads)
+
+ def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
+ o: torch.Tensor
+ q = q.view(-1, self.num_heads, self.head_dim)
+ k = k.view(-1, self.num_kv_heads, self.head_dim)
+ v = v.view(-1, self.num_kv_heads, self.head_dim)
+ context = get_context()
+ k_cache, v_cache = self.k_cache, self.v_cache
+
+ should_store_whole = (context.run_type == RunType.PREFILL)
+ if should_store_whole and k_cache.numel() and v_cache.numel():
+ store_kvcache(k, v, k_cache, v_cache, context.slot_mapping)
+
+ if context.run_type == RunType.PREFILL:
+ o = sparse_attn_varlen(q, k, v,
+ cu_seqlens_q=context.cu_seqlens_q,
+ cu_seqlens_k=context.cu_seqlens_k,
+ staircase_size=context.block_length)
+ else:
+ q = q.view(-1, context.block_length, self.num_heads, self.head_dim)
+ k = k.view(-1, context.block_length, self.num_kv_heads, self.head_dim)
+ v = v.view(-1, context.block_length, self.num_kv_heads, self.head_dim)
+ o = flash_attn_with_kvcache(q, k_cache=k_cache, v_cache=v_cache, k=k, v=v,
+ cache_seqlens=context.context_lens,
+ block_table=context.block_tables,
+ causal=False) # Assuming non-causal for benchmark consistency
+ o = o.view(-1, self.num_heads * self.head_dim)
+ return o
+
+
diff --git a/sample/jetengine_ext/layers/embed_head.py b/sample/jetengine_ext/layers/embed_head.py
new file mode 100644
index 0000000000000000000000000000000000000000..96968a29911e032c59932803792439840872f66b
--- /dev/null
+++ b/sample/jetengine_ext/layers/embed_head.py
@@ -0,0 +1,72 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+import torch.distributed as dist
+
+from jetengine_ext.engine.sequence import RunType
+from jetengine_ext.utils.context import get_context
+
+
+class VocabParallelEmbedding(nn.Module):
+
+ def __init__(
+ self,
+ num_embeddings: int,
+ embedding_dim: int,
+ ):
+ super().__init__()
+ self.tp_rank = dist.get_rank()
+ self.tp_size = dist.get_world_size()
+ assert num_embeddings % self.tp_size == 0
+ self.num_embeddings = num_embeddings
+ self.num_embeddings_per_partition = self.num_embeddings // self.tp_size
+ self.vocab_start_idx = self.num_embeddings_per_partition * self.tp_rank
+ self.vocab_end_idx = self.vocab_start_idx + self.num_embeddings_per_partition
+ self.weight = nn.Parameter(torch.empty(self.num_embeddings_per_partition, embedding_dim))
+ self.weight.weight_loader = self.weight_loader
+
+ def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
+ param_data = param.data
+ shard_size = param_data.size(0)
+ start_idx = self.tp_rank * shard_size
+ loaded_weight = loaded_weight.narrow(0, start_idx, shard_size)
+ assert param_data.size() == loaded_weight.size()
+ param_data.copy_(loaded_weight)
+
+ def forward(self, x: torch.Tensor):
+ if self.tp_size > 1:
+ mask = (x >= self.vocab_start_idx) & (x < self.vocab_end_idx)
+ x = mask * (x - self.vocab_start_idx)
+ y = F.embedding(x, self.weight)
+ if self.tp_size > 1:
+ y = mask.unsqueeze(1) * y
+ dist.all_reduce(y)
+ return y
+
+
+class ParallelLMHead(VocabParallelEmbedding):
+
+ def __init__(
+ self,
+ num_embeddings: int,
+ embedding_dim: int,
+ bias: bool = False,
+ ):
+ super().__init__(num_embeddings, embedding_dim)
+ if bias:
+ self.bias = nn.Parameter(torch.empty(self.num_embeddings_per_partition))
+ self.bias.weight_loader = self.weight_loader
+ else:
+ self.register_parameter("bias", None)
+
+ def forward(self, x: torch.Tensor):
+ context = get_context()
+ if context.run_type == RunType.PREFILL:
+ last_indices = context.cu_seqlens_q[1:] - 1
+ x = x[last_indices].contiguous()
+ logits = F.linear(x, self.weight, self.bias)
+ if self.tp_size > 1:
+ all_logits = [torch.empty_like(logits) for _ in range(self.tp_size)] if self.tp_rank == 0 else None
+ dist.gather(logits, all_logits, 0)
+ logits = torch.cat(all_logits, -1) if self.tp_rank == 0 else None
+ return logits
diff --git a/sample/jetengine_ext/layers/linear.py b/sample/jetengine_ext/layers/linear.py
new file mode 100644
index 0000000000000000000000000000000000000000..c625a9ebd0be699b869184ede34691d640367351
--- /dev/null
+++ b/sample/jetengine_ext/layers/linear.py
@@ -0,0 +1,174 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+import torch.distributed as dist
+
+
+def divide(numerator, denominator):
+ assert numerator % denominator == 0
+ return numerator // denominator
+
+
+class LinearBase(nn.Module):
+
+ def __init__(
+ self,
+ input_size: int,
+ output_size: int,
+ tp_dim: int | None = None,
+ ):
+ super().__init__()
+ self.input_size = input_size
+ self.output_size = output_size
+ self.tp_dim = tp_dim
+ self.tp_rank = dist.get_rank()
+ self.tp_size = dist.get_world_size()
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ raise NotImplementedError
+
+
+class ReplicatedLinear(LinearBase):
+
+ def __init__(
+ self,
+ input_size: int,
+ output_size: int,
+ bias: bool = False,
+ ):
+ super().__init__(input_size, output_size)
+ self.weight = nn.Parameter(torch.empty(self.output_size, self.input_size))
+ self.weight.weight_loader = self.weight_loader
+ if bias:
+ self.bias = nn.Parameter(torch.empty(self.output_size))
+ self.bias.weight_loader = self.weight_loader
+ else:
+ self.register_parameter("bias", None)
+
+ def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
+ param.data.copy_(loaded_weight)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return F.linear(x, self.weight, self.bias)
+
+
+class ColumnParallelLinear(LinearBase):
+
+ def __init__(
+ self,
+ input_size: int,
+ output_size: int,
+ bias: bool = False,
+ ):
+ super().__init__(input_size, output_size, 0)
+ self.input_size_per_partition = input_size
+ self.output_size_per_partition = divide(output_size, self.tp_size)
+
+ self.weight = nn.Parameter(torch.empty(self.output_size_per_partition, self.input_size))
+ self.weight.weight_loader = self.weight_loader
+ if bias:
+ self.bias = nn.Parameter(torch.empty(self.output_size_per_partition))
+ self.bias.weight_loader = self.weight_loader
+ else:
+ self.register_parameter("bias", None)
+
+ def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
+ param_data = param.data
+ shard_size = param_data.size(self.tp_dim)
+ start_idx = self.tp_rank * shard_size
+ loaded_weight = loaded_weight.narrow(self.tp_dim, start_idx, shard_size)
+ param_data.copy_(loaded_weight)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return F.linear(x, self.weight, self.bias)
+
+
+class MergedColumnParallelLinear(ColumnParallelLinear):
+
+ def __init__(
+ self,
+ input_size: int,
+ output_sizes: list[int],
+ bias: bool = False,
+ ):
+ self.output_sizes = output_sizes
+ super().__init__(input_size, sum(output_sizes), bias=bias)
+
+ def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: int):
+ param_data = param.data
+ shard_offset = sum(self.output_sizes[:loaded_shard_id]) // self.tp_size
+ shard_size = self.output_sizes[loaded_shard_id] // self.tp_size
+ param_data = param_data.narrow(self.tp_dim, shard_offset, shard_size)
+ loaded_weight = loaded_weight.chunk(self.tp_size, self.tp_dim)[self.tp_rank]
+ param_data.copy_(loaded_weight)
+
+
+class QKVParallelLinear(ColumnParallelLinear):
+
+ def __init__(
+ self,
+ hidden_size: int,
+ head_size: int,
+ total_num_heads: int,
+ total_num_kv_heads: int | None = None,
+ bias: bool = False,
+ ):
+ self.head_size = head_size
+ self.total_num_heads = total_num_heads
+ self.total_num_kv_heads = total_num_kv_heads or total_num_heads
+ tp_size = dist.get_world_size()
+ self.num_heads = divide(self.total_num_heads, tp_size)
+ self.num_kv_heads = divide(self.total_num_kv_heads, tp_size)
+ input_size = hidden_size
+ output_size = (self.total_num_heads + 2 * self.total_num_kv_heads) * self.head_size
+ super().__init__(input_size, output_size, bias)
+
+ def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str):
+ param_data = param.data
+ assert loaded_shard_id in ["q", "k", "v"]
+ if loaded_shard_id == "q":
+ shard_size = self.num_heads * self.head_size
+ shard_offset = 0
+ elif loaded_shard_id == "k":
+ shard_size = self.num_kv_heads * self.head_size
+ shard_offset = self.num_heads * self.head_size
+ else:
+ shard_size = self.num_kv_heads * self.head_size
+ shard_offset = self.num_heads * self.head_size + self.num_kv_heads * self.head_size
+ param_data = param_data.narrow(self.tp_dim, shard_offset, shard_size)
+ loaded_weight = loaded_weight.chunk(self.tp_size, self.tp_dim)[self.tp_rank]
+ param_data.copy_(loaded_weight)
+
+
+class RowParallelLinear(LinearBase):
+
+ def __init__(
+ self,
+ input_size: int,
+ output_size: int,
+ bias: bool = False,
+ ):
+ super().__init__(input_size, output_size, 1)
+ self.input_size_per_partition = divide(input_size, self.tp_size)
+ self.output_size_per_partition = output_size
+
+ self.weight = nn.Parameter(torch.empty(self.output_size, self.input_size_per_partition))
+ self.weight.weight_loader = self.weight_loader
+ if bias:
+ self.bias = nn.Parameter(torch.empty(self.output_size))
+ self.bias.weight_loader = self.weight_loader
+ else:
+ self.register_parameter("bias", None)
+
+ def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
+ param_data = param.data
+ shard_size = param_data.size(self.tp_dim)
+ start_idx = self.tp_rank * shard_size
+ loaded_weight = loaded_weight.narrow(self.tp_dim, start_idx, shard_size)
+ param_data.copy_(loaded_weight)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ y = F.linear(x, self.weight, self.bias if self.tp_rank == 0 else None)
+ if self.tp_size > 1:
+ dist.all_reduce(y)
+ return y
diff --git a/sample/jetengine_ext/layers/rotary_embedding.py b/sample/jetengine_ext/layers/rotary_embedding.py
new file mode 100644
index 0000000000000000000000000000000000000000..c473420f560204ce2a49b208754c78cd6a518c0a
--- /dev/null
+++ b/sample/jetengine_ext/layers/rotary_embedding.py
@@ -0,0 +1,68 @@
+from functools import lru_cache
+import torch
+from torch import nn
+
+
+def apply_rotary_emb(
+ x: torch.Tensor,
+ cos: torch.Tensor,
+ sin: torch.Tensor,
+) -> torch.Tensor:
+ cos = cos.unsqueeze(-2)
+ sin = sin.unsqueeze(-2)
+ x1, x2 = torch.chunk(x.to(torch.float32), 2, dim=-1)
+ y1 = x1 * cos - x2 * sin
+ y2 = x2 * cos + x1 * sin
+ return torch.cat((y1, y2), dim=-1).to(x.dtype)
+
+
+class RotaryEmbedding(nn.Module):
+
+ def __init__(
+ self,
+ head_size: int,
+ rotary_dim: int,
+ max_position_embeddings: int,
+ base: float,
+ ) -> None:
+ super().__init__()
+ self.head_size = head_size
+ assert rotary_dim == head_size
+ inv_freq = 1.0 / (base**(torch.arange(0, rotary_dim, 2, dtype=torch.float) / rotary_dim))
+ t = torch.arange(max_position_embeddings, dtype=torch.float)
+ freqs = torch.einsum("i,j -> ij", t, inv_freq)
+ cos = freqs.cos()
+ sin = freqs.sin()
+ cache = torch.cat((cos, sin), dim=-1)
+ self.register_buffer("cos_sin_cache", cache, persistent=False)
+
+ @torch.compile
+ def forward(
+ self,
+ positions: torch.Tensor,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ num_tokens = positions.size(0)
+ cos_sin = self.cos_sin_cache[positions]
+ cos, sin = cos_sin.chunk(2, dim=-1)
+ query_shape = query.shape
+ query = query.view(num_tokens, -1, self.head_size)
+ query = apply_rotary_emb(query, cos, sin).view(query_shape)
+ key_shape = key.shape
+ key = key.view(num_tokens, -1, self.head_size)
+ key = apply_rotary_emb(key, cos, sin).view(key_shape)
+ return query, key
+
+
+@lru_cache(1)
+def get_rope(
+ head_size: int,
+ rotary_dim: int,
+ max_position: int,
+ base: float,
+ rope_scaling: dict | None = None,
+):
+ assert rope_scaling is None
+ rotary_emb = RotaryEmbedding(head_size, rotary_dim, max_position, base)
+ return rotary_emb
diff --git a/sample/jetengine_ext/layers/sampler.py b/sample/jetengine_ext/layers/sampler.py
new file mode 100644
index 0000000000000000000000000000000000000000..b4b56dc5385a36796551726adc11500bd1fc8262
--- /dev/null
+++ b/sample/jetengine_ext/layers/sampler.py
@@ -0,0 +1,41 @@
+import torch
+from torch.nn import functional as F
+
+
+def top_k_logits(logits, k):
+ if k <= 0:
+ return logits
+ else:
+ values, _ = torch.topk(logits, k)
+ min_values = values[..., -1, None]
+ return torch.where(logits < min_values, torch.full_like(logits, float('-inf')), logits)
+
+def top_p_logits(logits, p):
+ sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+ cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+ sorted_mask = cumulative_probs > p
+ sorted_mask[..., 1:] = sorted_mask[..., :-1].clone()
+ sorted_mask[..., 0] = False
+ mask_indices = torch.scatter(torch.full_like(logits, False, dtype=torch.bool),
+ -1, sorted_indices, sorted_mask)
+ logits = logits.masked_fill(mask_indices, float('-inf'))
+ return logits
+
+def sample_with_temperature_topk_topp(logits, temperature=1.0, top_k=0, top_p=1.0):
+ orig_shape = logits.shape[:-1] # [batch, block]
+ vocab_size = logits.shape[-1]
+
+ logits = logits.reshape(-1, vocab_size) # [batch*block, vocab]
+
+ if temperature != 1.0:
+ logits = logits / temperature
+ if top_k > 0:
+ logits = top_k_logits(logits, top_k)
+ if top_p < 1.0:
+ logits = top_p_logits(logits, top_p)
+ probs = F.softmax(logits, dim=-1) # shape: [batch*block, vocab]
+ assert probs.dim() == 2
+ token = torch.multinomial(probs, num_samples=1) # [batch*block, 1]
+ token_prob = torch.gather(probs, -1, token) # [batch*block, 1]
+
+ return token.view(*orig_shape), token_prob.view(*orig_shape)
diff --git a/sample/jetengine_ext/llm.py b/sample/jetengine_ext/llm.py
new file mode 100644
index 0000000000000000000000000000000000000000..a650f582a499ab0cec4dcfcbfe8edb8b1d47aea1
--- /dev/null
+++ b/sample/jetengine_ext/llm.py
@@ -0,0 +1,5 @@
+from jetengine_ext.engine.llm_engine import LLMEngine
+
+
+class LLM(LLMEngine):
+ pass
diff --git a/sample/jetengine_ext/sampling_params.py b/sample/jetengine_ext/sampling_params.py
new file mode 100644
index 0000000000000000000000000000000000000000..97b30b64fc69ad7231fac44c5470be9c0f790b55
--- /dev/null
+++ b/sample/jetengine_ext/sampling_params.py
@@ -0,0 +1,18 @@
+from dataclasses import dataclass
+from typing import Literal
+
+@dataclass
+class SamplingParams:
+ temperature: float = 1.0
+ max_tokens: int = 64
+ ignore_eos: bool = False
+
+ # Block Diffusion Parameters
+ block_length: int = 4
+ denoising_steps: int = 4
+ dynamic_threshold: float = 0.9
+ eb_threshold: float = 0.35
+ topk: int = 0
+ topp: float = 1
+ remasking_strategy: Literal['sequential', 'low_confidence_static', 'low_confidence_dynamic', 'entropy_bounded'] = 'low_confidence_static'
+ stop_words: list[int] | None = None
diff --git a/sample/lavida_rl_rollout_video_reversal.py b/sample/lavida_rl_rollout_video_reversal.py
new file mode 100644
index 0000000000000000000000000000000000000000..c084de6da3bec9f89adad35cba85e67ae56f0be3
--- /dev/null
+++ b/sample/lavida_rl_rollout_video_reversal.py
@@ -0,0 +1,676 @@
+from __future__ import annotations
+import os
+import sys
+REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+LAVIDA_ROOT = os.path.join(REPO_ROOT, "LaViDa")
+sys.path.insert(0, LAVIDA_ROOT)
+sys.path.insert(1, REPO_ROOT)
+import math, json, os, time, re
+from dataclasses import dataclass
+from typing import List, Tuple
+from collections import defaultdict
+import numpy as np
+import torch
+from termcolor import cprint
+import torch.nn.functional as F
+from transformers import AutoTokenizer
+from llava.model.builder import load_pretrained_model
+from llava.mm_utils import process_images, tokenizer_image_token, process_video_with_decord
+from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
+from llava.conversation import conv_templates
+import multiprocessing as mp
+
+from omegaconf import DictConfig, ListConfig, OmegaConf
+
+from transformers.utils import logging as hf_logging
+hf_logging.set_verbosity_error()
+
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning)
+warnings.filterwarnings("ignore", category=FutureWarning)
+
+
+def get_config():
+ cli_conf = OmegaConf.from_cli()
+ yaml_conf = OmegaConf.load(cli_conf.config)
+ conf = OmegaConf.merge(yaml_conf, cli_conf)
+ return conf
+
+def add_gumbel_noise(logits, temperature):
+ if temperature == 0:
+ return logits
+ logits = logits.to(torch.float64)
+ noise = torch.rand_like(logits, dtype=torch.float64)
+ noise = (- torch.log(noise)) ** temperature
+ return logits.exp() / noise
+
+def get_num_transfer_tokens(mask_index, steps):
+ mask_num = mask_index.sum(dim=1, keepdim=True)
+ base = mask_num // steps
+ remainder = mask_num % steps
+ num_transfer_tokens = torch.zeros(mask_num.size(0), steps, device=mask_index.device, dtype=torch.int64) + base
+ for i in range(mask_num.size(0)):
+ num_transfer_tokens[i, :remainder[i]] += 1
+ return num_transfer_tokens
+
+@dataclass
+class DiffusionOutput:
+ sequences: torch.Tensor # final result (B, L_total) (GPU)
+ history: List[torch.Tensor] # all intermediate x (CPU)
+ nfe: int
+ L0: int # length of prompt
+
+def build_lavida_prompt(question: str, conv_template: str = "llada"):
+ import copy
+ conv = copy.deepcopy(conv_templates[conv_template])
+ conv.append_message(conv.roles[0], DEFAULT_IMAGE_TOKEN + "\n" + question)
+ conv.append_message(conv.roles[1], None)
+ return conv.get_prompt()
+
+@torch.no_grad()
+def generate_llada_v_with_history(
+ model, tokenizer, input_ids, images, image_sizes, modalities,
+ steps, gen_length, block_length, temperature,
+ target, unmask_threshold=None, mask_id=126336
+) -> DiffusionOutput:
+ position_ids = None
+ attention_mask = None
+ inputs_embeds = None
+
+ (_input_ids, position_ids, attention_mask, _pkv, inputs_embeds, _labels) = model.prepare_inputs_labels_for_multimodal(
+ input_ids, position_ids, attention_mask, None, None, images, modalities, image_sizes=image_sizes
+ )
+ assert input_ids.shape[1] < inputs_embeds.shape[1], "inputs_embeds must be longer than input_ids, because it contains image embeddings."
+ device = inputs_embeds.device
+ B = inputs_embeds.shape[0]
+ L0 = inputs_embeds.shape[1]
+
+ assert B == 1, "LLaDA-V only support batch size 1 for now"
+
+ total_len = L0 + gen_length
+ masked_embed = model.get_model().embed_tokens(torch.tensor([mask_id], device=device)) # (1, d)
+ x_embeds = masked_embed.repeat(B, total_len, 1)
+ x_embeds[:, :L0] = inputs_embeds
+ x_tokens = torch.full((B, total_len), mask_id, dtype=torch.long, device=device)
+ prompt_index = torch.zeros_like(x_tokens, dtype=torch.bool)
+ prompt_index[:, :L0] = True
+
+ assert gen_length % block_length == 0
+ num_blocks = gen_length // block_length
+ base, rem = divmod(steps, num_blocks)
+ steps_per_block = [base + (i < rem) for i in range(num_blocks)]
+
+ nfe = 0
+ hist: List[torch.Tensor] = []
+
+ for blk in range(num_blocks):
+ s = L0 + blk * block_length
+ e = L0 + (blk + 1) * block_length
+
+ cur_steps = steps_per_block[blk]
+
+ outputs = model.get_model()(
+ input_ids=None,
+ input_embeddings=x_embeds,
+ attention_mask=None,
+ use_cache=False,
+ )
+ logits = outputs.logits.float()
+
+ blk_tokens = x_tokens[:, s:e]
+ eos_mask = (blk_tokens == 126348)
+ if eos_mask.any():
+ first_eos_pos = torch.where(eos_mask[0])[0][0].item()
+ valid_region = torch.zeros_like(blk_tokens, dtype=torch.bool)
+ valid_region[:, :first_eos_pos] = True
+ else:
+ valid_region = torch.ones_like(blk_tokens, dtype=torch.bool)
+
+ mask_blk = (blk_tokens == mask_id) & valid_region
+ num_transfer = get_num_transfer_tokens(mask_blk, cur_steps)
+
+ x0_blk, tr_idx_blk = get_transfer_index(
+ logits[:, s:e], temperature, target,
+ mask_blk, blk_tokens, num_transfer[:, 0],
+ threshold=unmask_threshold
+ )
+ x_slice = x_tokens[:, s:e]
+ x_slice[tr_idx_blk] = x0_blk[tr_idx_blk]
+ x_tokens[:, s:e] = x_slice
+ x_embeds[:, s:e] = model.get_model().embed_tokens(x_tokens[:, s:e])
+ hist.append(x_tokens.clone().cpu())
+ nfe += 1
+
+ i = 1
+ while i < cur_steps:
+ nfe += 1
+ outputs = model.get_model()(
+ input_ids=None,
+ input_embeddings=x_embeds,
+ attention_mask=None,
+ use_cache=False,
+ )
+ logits = outputs.logits.float()
+
+ blk_tokens = x_tokens[:, s:e]
+ eos_mask = (blk_tokens == 126348)
+ if eos_mask.any():
+ first_eos_pos = torch.where(eos_mask[0])[0][0].item()
+ valid_region = torch.zeros_like(blk_tokens, dtype=torch.bool)
+ valid_region[:, :first_eos_pos] = True
+ else:
+ valid_region = torch.ones_like(blk_tokens, dtype=torch.bool)
+
+ mask_blk = (x_tokens[:, s:] == mask_id)
+ mask_blk[:, block_length:] = False
+ mask_blk[:, :block_length] = mask_blk[:, :block_length] & valid_region
+
+ x0, tr_idx = get_transfer_index(
+ logits[:, s:], temperature, target,
+ mask_blk, x_tokens[:, s:], num_transfer[:, i],
+ threshold=unmask_threshold
+ )
+
+ if tr_idx.any():
+ x_slice = x_tokens[:, s:]
+ x_slice[tr_idx] = x0[tr_idx]
+ x_tokens[:, s:] = x_slice
+ x0_embeds = model.get_model().embed_tokens(torch.where(tr_idx, x0, x_slice))
+ x_embeds[:, s:][tr_idx] = x0_embeds[tr_idx]
+
+ hist.append(x_tokens.clone().cpu())
+
+ remaining_masks = (x_tokens[:, s:e] == mask_id) & valid_region
+ if remaining_masks.sum() == 0:
+ break
+ i += 1
+
+ if (x_tokens[:, s:e] == 126348).any():
+ break
+
+ return DiffusionOutput(sequences=x_tokens, history=hist, nfe=nfe, L0=L0)
+
+def get_transfer_index(logits, temperature, target, mask_index, x, num_transfer_tokens, threshold=None):
+ logits_with_noise = add_gumbel_noise(logits, temperature=temperature)
+ logits_with_noise[..., 126336] = -float('inf')
+ x0 = torch.argmax(logits_with_noise, dim=-1)
+
+ if target == 'confidence':
+ p = F.softmax(logits.to(torch.float64), dim=-1)
+ x0_p = torch.squeeze(torch.gather(p, dim=-1, index=torch.unsqueeze(x0, -1)), -1)
+ elif target == 'margin_confidence':
+ p = F.softmax(logits.to(torch.float64), dim=-1)
+ top2 = torch.topk(p, 2, dim=-1).values
+ x0_p = top2[..., 0] - top2[..., 1]
+ elif target == 'neg_entropy':
+ p = F.softmax(logits.to(torch.float64), dim=-1)
+ x0_p = -torch.sum(p * torch.log(p + 1e-10), dim=-1)
+ elif target == 'random':
+ x0_p = torch.rand((x0.shape[0], x0.shape[1]), device=x0.device)
+ else:
+ raise NotImplementedError(target)
+
+ x0 = torch.where(mask_index, x0, x)
+
+ if threshold is not None:
+ selected = mask_index & (x0_p >= threshold)
+ for j in range(selected.shape[0]):
+ k = int(num_transfer_tokens[j].item() if torch.is_tensor(num_transfer_tokens[j]) else num_transfer_tokens[j])
+ if k <= 0: continue
+ num_selected = selected[j].sum().item()
+ if num_selected < k:
+ masked_scores = x0_p[j].masked_fill(~mask_index[j], float("-inf"))
+ _, top_indices = torch.topk(masked_scores, k=k)
+ selected[j, :] = False
+ selected[j, top_indices] = True
+ return x0, selected
+
+ confidence = x0_p.masked_fill(~mask_index, float("-inf"))
+ transfer_index = torch.zeros_like(x0, dtype=torch.bool, device=x0.device)
+ for j in range(confidence.shape[0]):
+ k = int(num_transfer_tokens[j].item() if torch.is_tensor(num_transfer_tokens[j]) else num_transfer_tokens[j])
+ if k <= 0: continue
+ _, sel = torch.topk(confidence[j], k=k)
+ transfer_index[j, sel] = True
+ return x0, transfer_index
+
+import random
+def random_select(data_list, random_k):
+ data_list = random.sample(data_list, random_k)
+ return data_list
+
+def get_prompt(data_i):
+ sys_prompt = 'Carefully watch the video and pay attention to the cause and sequence of events, the detail and movement of objects, and the action and pose of persons. Based on your observations, select the best option that accurately addresses the question.\n'
+ question_prompt="\nOnly give the best option."
+ multiple_choice_prompt = " Please provide only the single option letter (e.g., A, B, C, D, etc.) within the
+
+
+
+
+
+
+
+
+
+
+
+ Unmask Viewer
+ +
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/sample/trado_rl_rollout.py b/sample/trado_rl_rollout.py
new file mode 100644
index 0000000000000000000000000000000000000000..4d2ad4bd079d2c231f706994a4a7c63df8b6c181
--- /dev/null
+++ b/sample/trado_rl_rollout.py
@@ -0,0 +1,676 @@
+import os as _os
+_os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID")
+
+
+
+
+# Consolidate all caches into the local high-speed disk (NVMe or /dev/shm)
+# Local high-speed cache (NVMe or /dev/shm)
+_cache_root = "/dev/shm/torch_cache" # or "/local_nvme/torch_cache"
+_os.makedirs(_cache_root, exist_ok=True)
+_os.environ["TORCH_EXTENSIONS_DIR"] = _os.path.join(_cache_root, "torch_extensions")
+_os.environ["TRITON_CACHE_DIR"] = _os.path.join(_cache_root, "triton")
+_os.environ["XDG_CACHE_HOME"] = _cache_root
+_os.environ.setdefault("CUDA_MODULE_LOADING", "LAZY")
+
+
+
+_os.environ.setdefault("TORCH_NCCL_BLOCKING_WAIT", "1")
+_os.environ.setdefault("TORCH_NCCL_ASYNC_ERROR_HANDLING", "1")
+_os.environ.pop("NCCL_BLOCKING_WAIT", None)
+_os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
+
+import os
+import re
+import json
+from termcolor import cprint
+import random
+import torch.multiprocessing as mp
+from jinja2 import Template
+
+from omegaconf import DictConfig, ListConfig, OmegaConf
+def get_config():
+ cli_conf = OmegaConf.from_cli()
+ yaml_conf = OmegaConf.load(cli_conf.config)
+ conf = OmegaConf.merge(yaml_conf, cli_conf)
+ return conf
+
+
+
+# obtain prompt
+def get_prompt(data_i):
+ return Template(system_prompts).render(problem = data_i["question"])
+
+def extract_final_boxed_answer(s: str):
+ tag = r'\boxed{'
+ start = s.rfind(tag) # last \boxed{
+ if start == -1:
+ return "Can not extract the answer!"
+
+ i = start + len(tag)
+ depth = 1 # we are already inside one '{'
+ buf = []
+
+ while i < len(s) and depth:
+ ch = s[i]
+ if ch == '{':
+ depth += 1
+ elif ch == '}':
+ depth -= 1
+ if depth == 0: # matching '}' for the opening \boxed{
+ break
+ buf.append(ch)
+ i += 1
+
+ return ''.join(buf) if depth == 0 else "Can not extract the answer!"
+
+
+
+
+def extract_code(full_output):
+ matches = re.findall(r"```python(.*?)```", full_output, re.DOTALL)
+ if matches:
+ code_output = matches[-1].strip()
+ else:
+ code_output = "We can not extract the code in the output. "
+ return code_output
+
+
+def get_data_chunk(data, num_node, node_idx):
+ total = len(data)
+ chunk_size = (total + num_node - 1) // num_node
+ start_idx = node_idx * chunk_size
+ end_idx = min((node_idx + 1) * chunk_size, total)
+ return data[start_idx:end_idx]
+
+
+import socket
+
+def _patch_safe_destroy():
+ import torch.distributed as dist
+ _real_destroy = dist.destroy_process_group
+ def _safe_destroy(group=None):
+ try:
+ if not dist.is_available():
+ return
+ try:
+ if not dist.is_initialized():
+ return
+ except Exception:
+ return
+ _real_destroy(group)
+ except AssertionError:
+ pass
+ dist.destroy_process_group = _safe_destroy
+
+
+
+def _llm_worker_run(args):
+ (model_path, tp, block_size, sampling_kwargs, vis_ids,
+ prompts_slice, indices_slice, enforce_eager, max_active, store_port) = args
+
+ import os
+ # 1) Setup environment (critical for correct worker behavior)
+ os.environ.setdefault("TORCH_NCCL_BLOCKING_WAIT", "1")
+ os.environ.setdefault("TORCH_NCCL_ASYNC_ERROR_HANDLING", "1")
+ os.environ.pop("NCCL_BLOCKING_WAIT", None)
+ os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
+ os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, vis_ids))
+ #os.environ.setdefault("TORCH_EXTENSIONS_DIR", f"/tmp/torch_ext_worker_{store_port}")
+ os.environ["MASTER_ADDR"] = "127.0.0.1"
+ os.environ["MASTER_PORT"] = str(store_port)
+ os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+ # 1.1) Create a per-worker `sitecustomize.py` and inject it into PYTHONPATH
+ # (this must be done before importing torch/jetengine)
+ patch_dir = f"/tmp/je_site_{store_port}"
+ os.makedirs(patch_dir, exist_ok=True)
+ patch_file = os.path.join(patch_dir, "sitecustomize.py")
+ # Important: the content must start at column 0 (no indentation)!
+ with open(patch_file, "w") as _f:
+ _f.write(
+ "import os\n"
+ "import torch.distributed as dist\n"
+ "_real = dist.init_process_group\n"
+ "def _wrapped(backend, init_method=None, *args, **kwargs):\n"
+ " port = os.environ.get('JE_TCP_PORT')\n"
+ " if port and isinstance(init_method, str) and init_method.startswith('tcp://localhost:2333'):\n"
+ " init_method = f'tcp://127.0.0.1:{port}'\n"
+ " return _real(backend, init_method, *args, **kwargs)\n"
+ "dist.init_process_group = _wrapped\n"
+ )
+ os.environ["PYTHONPATH"] = patch_dir + (":" + os.environ["PYTHONPATH"] if "PYTHONPATH" in os.environ else "")
+ os.environ["JE_TCP_PORT"] = str(store_port)
+
+ # 2) Import torch and patch the current worker process
+ import torch
+ import torch.distributed as dist
+ _patch_dist_port(store_port) # Patch port binding for this process
+ _patch_safe_destroy() # Avoid AssertionError in destroy_process_group
+ torch.cuda.set_device(0) # From this worker’s perspective, cuda:0 is the first visible device
+
+ # For debugging: print the worker’s CUDA_VISIBLE_DEVICES and assigned port
+ print(f"[worker pid={os.getpid()}] CVD={os.environ['CUDA_VISIBLE_DEVICES']}, port={store_port}, prompts={len(prompts_slice)}", flush=True)
+
+ # 3) Import jetengine and create the engine
+ # (child processes inherit the sitecustomize patch)
+ from jetengine_ext.llm import LLM
+ from jetengine_ext.sampling_params import SamplingParams
+
+ llm = LLM(
+ model_path,
+ enforce_eager=enforce_eager,
+ tensor_parallel_size=tp,
+ mask_token_id=151669,
+ block_length=block_size
+ )
+ sp = SamplingParams(**sampling_kwargs)
+ outs = llm.generate_streaming(prompts_slice, sp, max_active=max_active)
+ #seq = [o["text"] for o in outs]
+ #print(outs[0]["first_unmask_times"])
+ triples = []
+ for j, o in enumerate(outs):
+ triples.append((
+ indices_slice[j], # Global index (used to restore original order)
+ o["text"], # Generated text
+ o.get("first_unmask_times", None) # Optional time series aligned with completion tokens
+ ))
+
+ try:
+ if hasattr(llm, "shutdown"):
+ llm.shutdown()
+ except Exception:
+ pass
+
+ return triples
+
+
+
+def _llm_worker_entry(args, out_q):
+ import traceback, os
+ try:
+ res = _llm_worker_run(args)
+ out_q.put(("ok", res))
+ except Exception:
+ out_q.put(("err", {
+ "pid": os.getpid(),
+ "port": args[-1], # store_port
+ "traceback": traceback.format_exc(),
+ }))
+
+
+def _find_free_port():
+ s = socket.socket(); s.bind(('', 0))
+ p = s.getsockname()[1]; s.close()
+ return p
+
+def _patch_dist_port(port: int):
+ import torch.distributed as _dist
+ _real_init = _dist.init_process_group
+
+ def _wrapped(backend, init_method=None, *args, **kwargs):
+ # jetengine internally hardcodes "tcp://localhost:2333" — replace the port here
+ if isinstance(init_method, str) and init_method.startswith("tcp://localhost:2333"):
+ init_method = f"tcp://127.0.0.1:{port}"
+ return _real_init(backend, init_method, *args, **kwargs)
+
+ _dist.init_process_group = _wrapped
+
+
+import random
+def random_select(data_list, random_k):
+ data_list = random.sample(data_list, random_k)
+ return data_list
+
+if __name__ == "__main__":
+
+
+
+ tp = int(get_config().rollout.tensor_parallel_size) # Or check after loading config
+
+ if tp == 1:
+ os.environ.setdefault("TORCH_NCCL_ASYNC_ERROR_HANDLING", "1")
+ os.environ.setdefault("TORCH_NCCL_BLOCKING_WAIT", "1")
+ # These two are NCCL’s own variables, keep using the NCCL_ prefix
+ os.environ.setdefault("NCCL_P2P_DISABLE", "1")
+ os.environ.setdefault("NCCL_IB_DISABLE", "1")
+ else:
+ # For multi-GPU communication, do not disable P2P/IB;
+ # also clean up related variables (both old and new names)
+ for k in [
+ "NCCL_P2P_DISABLE", "NCCL_IB_DISABLE",
+ "TORCH_NCCL_BLOCKING_WAIT", "TORCH_NCCL_ASYNC_ERROR_HANDLING",
+ "NCCL_BLOCKING_WAIT", "NCCL_ASYNC_ERROR_HANDLING",
+ ]:
+ os.environ.pop(k, None)
+
+
+
+ from transformers import AutoTokenizer
+
+ # --- graceful shutdown & unique port ---
+ import os, sys, atexit, signal, torch.distributed as dist
+
+
+ # 2) Automatically set compile architecture according to the local GPU
+ # (do NOT hardcode 8.0)
+ def _set_arch():
+ try:
+ if torch.cuda.is_available():
+ major, minor = torch.cuda.get_device_capability(0)
+ os.environ["TORCH_CUDA_ARCH_LIST"] = f"{major}.{minor}"
+ except Exception:
+ pass
+ _set_arch()
+
+
+
+ # 1) Use a new port at each startup to avoid conflicts with 2333
+ if "MASTER_PORT" not in os.environ:
+ os.environ["MASTER_ADDR"] = "127.0.0.1"
+ os.environ["MASTER_PORT"] = str(_find_free_port())
+ # (If JetEngine hardcodes tcp://localhost:2333 instead of using env://,
+ # see the “special case” section at the end)
+
+ # 2) Intercept Ctrl-C/TERM to destroy distributed groups & engine gracefully
+ _llm = None
+ _child_ps = [] # If you create your own mp.Process/Pool, append objects here
+
+ def _cleanup():
+ # 2.1) Shutdown JetEngine engine (if API available)
+ global _llm
+ try:
+ if _llm is not None and hasattr(_llm, "shutdown"):
+ _llm.shutdown()
+ except Exception:
+ pass
+ # 2.3) Kill/join child processes
+ for p in _child_ps:
+ try:
+ if hasattr(p, "terminate"): p.terminate()
+ except Exception:
+ pass
+ for p in _child_ps:
+ try:
+ if hasattr(p, "join"): p.join(timeout=2)
+ except Exception:
+ pass
+
+ atexit.register(_cleanup)
+ def _sig_handler(sig, frame):
+ _cleanup()
+ # 130: standard exit code for SIGINT; 143: for SIGTERM
+ sys.exit(130 if sig == signal.SIGINT else 143)
+
+ signal.signal(signal.SIGINT, _sig_handler)
+ signal.signal(signal.SIGTERM, _sig_handler)
+
+
+
+
+
+
+
+
+
+ config = get_config()
+
+ try:
+ if mp.get_start_method(allow_none=True) != "spawn":
+ mp.set_start_method("spawn", force=True)
+ except RuntimeError:
+ pass
+
+
+
+ system_prompts = '''<|im_start|>user\n{{problem}}\nPlease reason step by step, and put your final answer within \\boxed{}.<|im_end|>\n<|im_start|>assistant\n'''
+ if config.rollout.start_with_think:
+ system_prompts = '''<|im_start|>user\nYou need to put your final answer in \\boxed{}. This is the problem:\n{{problem}}<|im_end|>\n<|im_start|>assistant
+
+
+
+
+
+
+
+
+
+
+
+