trending/miles
1
0
Fork 0
mirror of https://github.com/radixark/miles.git synced 2026-10-02 07:14:53 +08:00
Code Issues Packages Projects Releases Wiki Activity

update

Browse Source
This commit is contained in:
Yusheng Su
2025-07-01 07:10:27 +00:00
parent 0f0a4388d6
commit b4e669f2ec
7 changed files with 327 additions and 1066 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+4
View File
@@ -170,6 +170,10 @@ For complete usage instructions, please refer to the [Usage Documentation](docs/
- For debugging tips, please refer to the [Debugging Guide](docs/en/debug.md)
## Hardware Support
- Nvidia: refer to this repo README
- AMD: refer to the [toturial](docs/en/amd_toturial.md)
## FAQ & Acknowledgements
- For frequently asked questions, please see the [Q\&A](docs/en/qa.md)
docker/Dockerfile.rocm
+5 -5
View File
@@ -216,11 +216,11 @@ WORKDIR /workspace/
#########################################
#### Install megatron-core###############
#########################################
RUN pip uninstall -y megatron-core && \
git clone https://github.com/yushengsu-thu/Megatron-LM-amd_version.git && \
cd Megatron-LM-amd_version && \
pip install -vvv -e . && \
cd /workspace/
RUN pip uninstall -y megatron-core
COPY Megatron-LM-amd_version /workspace/Megatron-LM-amd_version/
WORKDIR /workspace/Megatron-LM-amd_version
RUN pip install -vvv -e .
WORKDIR /workspace/
#########################################
#########################################
#########################################
docker/patch/megatron-sandwich_norm.patch
-136
View File
@@ -1,136 +0,0 @@
diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py
index 07fbd64f..cc8b8ec7 100755
--- a/megatron/core/models/gpt/gpt_layer_specs.py
+++ b/megatron/core/models/gpt/gpt_layer_specs.py
@@ -74,6 +74,8 @@ def get_gpt_layer_with_transformer_engine_spec(
qk_l2_norm: Optional[bool] = False,
use_te_op_fuser: Optional[bool] = False,
use_kitchen: bool = False,
+ post_self_attn_layernorm: bool = False,
+ post_mlp_layernorm: bool = False,
) -> ModuleSpec:
"""Use this spec to use lower-level Transformer Engine modules (required for fp8 training).
@@ -173,9 +175,11 @@ def get_gpt_layer_with_transformer_engine_spec(
),
),
self_attn_bda=get_bias_dropout_add,
+ post_self_attn_layernorm=TENorm if post_self_attn_layernorm else IdentityOp,
pre_mlp_layernorm=backend.layer_norm() if num_experts else IdentityOp,
mlp=mlp,
mlp_bda=get_bias_dropout_add,
+ post_mlp_layernorm=TENorm if post_mlp_layernorm else IdentityOp,
sharded_state_dict_keys_map={
'mlp.0.weight': 'mlp.linear_fc1.layer_norm_weight',
'mlp.0.bias': 'mlp.linear_fc1.layer_norm_bias',
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index a256d70f..0d436eb0 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -167,6 +167,9 @@ class TransformerConfig(ModelParallelConfig):
qk_layernorm: bool = False
"""Whether to apply `normalization` type of normalization to the query and key embeddings."""
+ post_self_attn_layernorm: bool = False
+ post_mlp_layernorm: bool = False
+
test_mode: bool = False
"""Whether to run real-time tests."""
diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py
index 51c90f5a..05c34955 100644
--- a/megatron/core/transformer/transformer_layer.py
+++ b/megatron/core/transformer/transformer_layer.py
@@ -228,6 +228,7 @@ class TransformerLayerSubmodules:
input_layernorm: Union[ModuleSpec, type] = IdentityOp
self_attention: Union[ModuleSpec, type] = IdentityOp
self_attn_bda: Union[ModuleSpec, type] = IdentityFuncOp
+ post_self_attn_layernorm: Union[ModuleSpec, type] = IdentityOp
pre_cross_attn_layernorm: Union[ModuleSpec, type] = IdentityOp
cross_attention: Union[ModuleSpec, type] = IdentityOp
@@ -236,6 +237,7 @@ class TransformerLayerSubmodules:
pre_mlp_layernorm: Union[ModuleSpec, type] = IdentityOp
mlp: Union[ModuleSpec, type] = IdentityOp
mlp_bda: Union[ModuleSpec, type] = IdentityFuncOp
+ post_mlp_layernorm: Union[ModuleSpec, type] = IdentityOp
# Mapping for sharded tensor keys to be applied in `sharded_state_dict` method
sharded_state_dict_keys_map: Dict[str, str] = field(default_factory=dict)
@@ -338,6 +340,14 @@ class TransformerLayer(MegatronModule, BaseTransformerLayer):
# [Module 3: BiasDropoutFusion]
self.self_attn_bda = build_module(submodules.self_attn_bda)
+ self.post_self_attn_layernorm = build_module(
+ submodules.post_self_attn_layernorm,
+ config=self.config,
+ hidden_size=self.config.hidden_size,
+ eps=self.config.layernorm_epsilon,
+ )
+
+
# [Module 4: Post SelfAttention] Optional Layernorm after self-attn
self.pre_cross_attn_layernorm = build_module(
submodules.pre_cross_attn_layernorm,
@@ -395,6 +405,13 @@ class TransformerLayer(MegatronModule, BaseTransformerLayer):
# [Module 9: BiasDropoutFusion]
self.mlp_bda = build_module(submodules.mlp_bda)
+ self.post_mlp_layernorm = build_module(
+ submodules.post_mlp_layernorm,
+ config=self.config,
+ hidden_size=self.config.hidden_size,
+ eps=self.config.layernorm_epsilon
+ )
+
self.recompute_input_layernorm = False
self.recompute_pre_mlp_layernorm = False
self.recompute_mlp = False
@@ -518,6 +535,11 @@ class TransformerLayer(MegatronModule, BaseTransformerLayer):
attention_output_with_bias[0]
)
+ attention_output, attention_output_bias = attention_output_with_bias
+ attention_output = self.post_self_attn_layernorm(attention_output)
+ attention_output_with_bias = (attention_output, attention_output_bias)
+
+
# TODO: could we move `bias_dropout_add_exec_handler` itself
# inside the module provided in the `bias_dropout_add_spec` module?
nvtx_range_push(suffix="self_attn_bda")
@@ -618,6 +640,10 @@ class TransformerLayer(MegatronModule, BaseTransformerLayer):
else:
mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output)
+ mlp_output, mlp_output_bias = mlp_output_with_bias
+ mlp_output = self.post_mlp_layernorm(mlp_output)
+ mlp_output_with_bias = (mlp_output, mlp_output_bias)
+
if self.recompute_pre_mlp_layernorm:
# discard the output of the pre-mlp layernorm and register the recompute
# as a gradient hook of mlp_output_with_bias[0]
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 7b1b35b1..c0c510e9 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -1158,6 +1158,9 @@ def core_transformer_config_from_args(args, config_class=None):
if args.is_hybrid_model:
kw_args['is_hybrid_model'] = args.is_hybrid_model
+ kw_args['post_self_attn_layernorm'] = args.post_self_attn_layernorm
+ kw_args['post_mlp_layernorm'] = args.post_mlp_layernorm
+
# handle quantization config
# NOTE: Kitchen arguments are only added to the namespace when
# Kitchen library is available.
@@ -1439,6 +1442,10 @@ def _add_network_size_args(parser):
action='store_true',
help='If set, use original BERT residula connection '
'ordering.')
+ group.add_argument('--post-self-attn-layernorm', action='store_true',
+ help='If set, use post self attention layernorm.')
+ group.add_argument('--post-mlp-layernorm', action='store_true',
+ help='If set, use post MLP layernorm.')
group.add_argument('--openai-gelu', action='store_true',
help='Use OpenAIs GeLU implementation. This option'
'should not be used unless for backward compatibility'
docker/patch/sglang.patch
-906
View File
@@ -1,906 +0,0 @@
diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py
index 0b641d34..bc7cba90 100644
--- a/python/sglang/srt/configs/model_config.py
+++ b/python/sglang/srt/configs/model_config.py
@@ -390,15 +390,15 @@ class ModelConfig:
if quant_cfg is not None:
quant_method = quant_cfg.get("quant_method", "").lower()
- # Detect which checkpoint is it
- for _, method in QUANTIZATION_METHODS.items():
- quantization_override = method.override_quantization_method(
- quant_cfg, self.quantization
- )
- if quantization_override:
- quant_method = quantization_override
- self.quantization = quantization_override
- break
+ # # Detect which checkpoint is it
+ # for _, method in QUANTIZATION_METHODS.items():
+ # quantization_override = method.override_quantization_method(
+ # quant_cfg, self.quantization
+ # )
+ # if quantization_override:
+ # quant_method = quantization_override
+ # self.quantization = quantization_override
+ # break
# Verify quantization configurations.
if self.quantization is None:
diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py
index 75bccc9d..0ca97353 100644
--- a/python/sglang/srt/entrypoints/engine.py
+++ b/python/sglang/srt/entrypoints/engine.py
@@ -403,12 +403,15 @@ class Engine(EngineBase):
self.tokenizer_manager.init_weights_update_group(obj, None)
)
- def update_weights_from_distributed(self, name: str, dtype, shape):
+ def update_weights_from_distributed(
+ self, names: str, dtypes, shapes, group_name: str
+ ):
"""Update weights from distributed source."""
obj = UpdateWeightsFromDistributedReqInput(
- name=name,
- dtype=dtype,
- shape=shape,
+ names=names,
+ dtypes=dtypes,
+ shapes=shapes,
+ group_name=group_name,
)
loop = asyncio.get_event_loop()
return loop.run_until_complete(
diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py
index 06c67343..672b294d 100644
--- a/python/sglang/srt/entrypoints/http_server.py
+++ b/python/sglang/srt/entrypoints/http_server.py
@@ -605,6 +605,20 @@ async def separate_reasoning_request(obj: SeparateReasoningReqInput, request: Re
return ORJSONResponse(content=response_data, status_code=200)
+@app.post("/pause_generation")
+async def pause_generation(request: Request):
+ """Pause generation."""
+ await _global_state.tokenizer_manager.pause_generation()
+ return ORJSONResponse(content="Generation paused.", status_code=200)
+
+
+@app.post("/continue_generation")
+async def continue_generation(request: Request):
+ """Continue generation."""
+ await _global_state.tokenizer_manager.continue_generation()
+ return ORJSONResponse(content="Generation continued.", status_code=200)
+
+
##### OpenAI-compatible API endpoints #####
@@ -806,104 +820,105 @@ def _wait_and_warmup(
image_token_text: str,
launch_callback: Optional[Callable[[], None]] = None,
):
- headers = {}
- url = server_args.url()
- if server_args.api_key:
- headers["Authorization"] = f"Bearer {server_args.api_key}"
+ if not server_args.skip_warmup:
+ headers = {}
+ url = server_args.url()
+ if server_args.api_key:
+ headers["Authorization"] = f"Bearer {server_args.api_key}"
+
+ # Wait until the server is launched
+ success = False
+ for _ in range(120):
+ time.sleep(1)
+ try:
+ res = requests.get(url + "/get_model_info", timeout=5, headers=headers)
+ assert res.status_code == 200, f"{res=}, {res.text=}"
+ success = True
+ break
+ except (AssertionError, requests.exceptions.RequestException):
+ last_traceback = get_exception_traceback()
+ pass
+
+ if not success:
+ if pipe_finish_writer is not None:
+ pipe_finish_writer.send(last_traceback)
+ logger.error(f"Initialization failed. warmup error: {last_traceback}")
+ kill_process_tree(os.getpid())
+ return
+
+ model_info = res.json()
+
+ # Send a warmup request
+ request_name = "/generate" if model_info["is_generation"] else "/encode"
+ max_new_tokens = 8 if model_info["is_generation"] else 1
+ json_data = {
+ "sampling_params": {
+ "temperature": 0,
+ "max_new_tokens": max_new_tokens,
+ },
+ }
+ if server_args.skip_tokenizer_init:
+ json_data["input_ids"] = [[10, 11, 12] for _ in range(server_args.dp_size)]
+ # TODO Workaround the bug that embedding errors for list of size 1
+ if server_args.dp_size == 1:
+ json_data["input_ids"] = json_data["input_ids"][0]
+ else:
+ json_data["text"] = ["The capital city of France is"] * server_args.dp_size
+ # TODO Workaround the bug that embedding errors for list of size 1
+ if server_args.dp_size == 1:
+ json_data["text"] = json_data["text"][0]
+
+ # Debug dumping
+ if server_args.debug_tensor_dump_input_file:
+ json_data.pop("text", None)
+ json_data["input_ids"] = np.load(
+ server_args.debug_tensor_dump_input_file
+ ).tolist()
+ json_data["sampling_params"]["max_new_tokens"] = 0
- # Wait until the server is launched
- success = False
- for _ in range(120):
- time.sleep(1)
try:
- res = requests.get(url + "/get_model_info", timeout=5, headers=headers)
- assert res.status_code == 200, f"{res=}, {res.text=}"
- success = True
- break
- except (AssertionError, requests.exceptions.RequestException):
+ if server_args.disaggregation_mode == "null":
+ res = requests.post(
+ url + request_name,
+ json=json_data,
+ headers=headers,
+ timeout=600,
+ )
+ assert res.status_code == 200, f"{res}"
+ else:
+ logger.info(f"Start of prefill warmup ...")
+ json_data = {
+ "sampling_params": {
+ "temperature": 0.0,
+ "max_new_tokens": 8,
+ "ignore_eos": True,
+ },
+ "bootstrap_host": [FakeBootstrapHost] * server_args.dp_size,
+ # This is a hack to ensure fake transfer is enabled during prefill warmup
+ # ensure each dp rank has a unique bootstrap_room during prefill warmup
+ "bootstrap_room": [
+ i * (2**63 // server_args.dp_size) + (i % server_args.tp_size)
+ for i in range(server_args.dp_size)
+ ],
+ "input_ids": [[0, 1, 2, 3]] * server_args.dp_size,
+ }
+ res = requests.post(
+ url + request_name,
+ json=json_data,
+ headers=headers,
+ timeout=1800, # because of deep gemm precache is very long if not precache.
+ )
+ logger.info(
+ f"End of prefill warmup with status {res.status_code}, resp: {res.json()}"
+ )
+
+ except Exception:
last_traceback = get_exception_traceback()
- pass
-
- if not success:
- if pipe_finish_writer is not None:
- pipe_finish_writer.send(last_traceback)
- logger.error(f"Initialization failed. warmup error: {last_traceback}")
- kill_process_tree(os.getpid())
- return
-
- model_info = res.json()
-
- # Send a warmup request
- request_name = "/generate" if model_info["is_generation"] else "/encode"
- max_new_tokens = 8 if model_info["is_generation"] else 1
- json_data = {
- "sampling_params": {
- "temperature": 0,
- "max_new_tokens": max_new_tokens,
- },
- }
- if server_args.skip_tokenizer_init:
- json_data["input_ids"] = [[10, 11, 12] for _ in range(server_args.dp_size)]
- # TODO Workaround the bug that embedding errors for list of size 1
- if server_args.dp_size == 1:
- json_data["input_ids"] = json_data["input_ids"][0]
- else:
- json_data["text"] = ["The capital city of France is"] * server_args.dp_size
- # TODO Workaround the bug that embedding errors for list of size 1
- if server_args.dp_size == 1:
- json_data["text"] = json_data["text"][0]
-
- # Debug dumping
- if server_args.debug_tensor_dump_input_file:
- json_data.pop("text", None)
- json_data["input_ids"] = np.load(
- server_args.debug_tensor_dump_input_file
- ).tolist()
- json_data["sampling_params"]["max_new_tokens"] = 0
-
- try:
- if server_args.disaggregation_mode == "null":
- res = requests.post(
- url + request_name,
- json=json_data,
- headers=headers,
- timeout=600,
- )
- assert res.status_code == 200, f"{res}"
- else:
- logger.info(f"Start of prefill warmup ...")
- json_data = {
- "sampling_params": {
- "temperature": 0.0,
- "max_new_tokens": 8,
- "ignore_eos": True,
- },
- "bootstrap_host": [FakeBootstrapHost] * server_args.dp_size,
- # This is a hack to ensure fake transfer is enabled during prefill warmup
- # ensure each dp rank has a unique bootstrap_room during prefill warmup
- "bootstrap_room": [
- i * (2**63 // server_args.dp_size) + (i % server_args.tp_size)
- for i in range(server_args.dp_size)
- ],
- "input_ids": [[0, 1, 2, 3]] * server_args.dp_size,
- }
- res = requests.post(
- url + request_name,
- json=json_data,
- headers=headers,
- timeout=1800, # because of deep gemm precache is very long if not precache.
- )
- logger.info(
- f"End of prefill warmup with status {res.status_code}, resp: {res.json()}"
- )
-
- except Exception:
- last_traceback = get_exception_traceback()
- if pipe_finish_writer is not None:
- pipe_finish_writer.send(last_traceback)
- logger.error(f"Initialization failed. warmup error: {last_traceback}")
- kill_process_tree(os.getpid())
- return
+ if pipe_finish_writer is not None:
+ pipe_finish_writer.send(last_traceback)
+ logger.error(f"Initialization failed. warmup error: {last_traceback}")
+ kill_process_tree(os.getpid())
+ return
# Debug print
# logger.info(f"{res.json()=}")
diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py
index c779f1f1..a3b1db7c 100644
--- a/python/sglang/srt/layers/quantization/fp8.py
+++ b/python/sglang/srt/layers/quantization/fp8.py
@@ -325,10 +325,10 @@ class Fp8LinearMethod(LinearMethodBase):
layer.input_scale = None
else:
weight, weight_scale = layer.weight.data, layer.weight_scale_inv.data
- layer.weight = torch.nn.Parameter(weight, requires_grad=False)
- layer.weight_scale_inv = torch.nn.Parameter(
- weight_scale, requires_grad=False
- )
+ # layer.weight = torch.nn.Parameter(weight, requires_grad=False)
+ # layer.weight_scale_inv = torch.nn.Parameter(
+ # weight_scale, requires_grad=False
+ # )
return
layer.weight = torch.nn.Parameter(layer.weight.data, requires_grad=False)
diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py
index 40c220c4..153be20a 100644
--- a/python/sglang/srt/managers/io_struct.py
+++ b/python/sglang/srt/managers/io_struct.py
@@ -727,9 +727,11 @@ class UpdateWeightFromDiskReqOutput:
@dataclass
class UpdateWeightsFromDistributedReqInput:
- name: str
- dtype: str
- shape: List[int]
+ names: List[str]
+ dtypes: List[str]
+ shapes: List[List[int]]
+ # The group name
+ group_name: str = "weight_update_group"
@dataclass
diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py
index a44515ab..532802df 100644
--- a/python/sglang/srt/managers/scheduler.py
+++ b/python/sglang/srt/managers/scheduler.py
@@ -2036,6 +2036,9 @@ class Scheduler(
# Delete requests in the waiting queue
to_del = []
for i, req in enumerate(self.waiting_queue):
+ if recv_req.rid == "":
+ to_del.append(i)
+ continue
if req.rid.startswith(recv_req.rid):
to_del.append(i)
@@ -2065,6 +2068,9 @@ class Scheduler(
reqs = self.running_batch.reqs + self.cur_batch.reqs
for req in reqs:
+ if recv_req.rid == "" and not req.finished():
+ req.to_abort = True
+ continue
if req.rid.startswith(recv_req.rid) and not req.finished():
# Abort method 3: set `to_abort=True`
# The request will still run one decode forward pass.
@@ -2096,7 +2102,7 @@ class Scheduler(
) -> Tuple[bool, str]:
"""Update the online model parameter."""
success, message = self.tp_worker.update_weights_from_distributed(recv_req)
- if success:
+ if success and False:
flash_cache_success = self.flush_cache()
assert flash_cache_success, "Cache flush failed after updating weights"
else:
@@ -2105,7 +2111,11 @@ class Scheduler(
def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput):
"""Update the online model parameter from tensors."""
- success, message = self.tp_worker.update_weights_from_tensor(recv_req)
+ if self.draft_worker is not None:
+ success, message = self.draft_worker.update_weights_from_tensor(recv_req)
+ else:
+ success, message = self.tp_worker.update_weights_from_tensor(recv_req)
+
# TODO extract common code b/t update_weights_from_distributed and update_weights_from_tensor later
if success:
if recv_req.flush_cache:
diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py
index d71bbdf0..f3968afb 100644
--- a/python/sglang/srt/managers/tokenizer_manager.py
+++ b/python/sglang/srt/managers/tokenizer_manager.py
@@ -199,6 +199,8 @@ class TokenizerManager:
self.is_image_gen = self.model_config.is_image_gen
self.context_len = self.model_config.context_len
self.image_token_id = self.model_config.image_token_id
+ self._updating = False
+ self._cond = asyncio.Condition()
if self.model_config.is_multimodal:
import_processors()
@@ -407,6 +409,8 @@ class TokenizerManager:
request: Optional[fastapi.Request] = None,
):
created_time = time.time()
+ async with self._cond:
+ await self._cond.wait_for(lambda: not self._updating)
self.auto_create_handle_loop()
@@ -789,7 +793,7 @@ class TokenizerManager:
return (await self.flush_cache_communicator(FlushCacheReqInput()))[0]
def abort_request(self, rid: str):
- if rid not in self.rid_to_state:
+ if rid != "" and rid not in self.rid_to_state:
return
req = AbortReq(rid)
self.send_to_scheduler.send_pyobj(req)
@@ -844,6 +848,16 @@ class TokenizerManager:
self.auto_create_handle_loop()
await self.expert_distribution_communicator(ExpertDistributionReq.DUMP_RECORD)
+ async def pause_generation(self):
+ async with self._cond:
+ self._updating = True
+ self.abort_request("")
+
+ async def continue_generation(self):
+ async with self._cond:
+ self._updating = False
+ self._cond.notify_all()
+
async def update_weights_from_disk(
self,
obj: UpdateWeightFromDiskReqInput,
@@ -1423,7 +1437,21 @@ class TokenizerManager:
asyncio.create_task(asyncio.to_thread(background_task))
def _handle_abort_req(self, recv_obj):
- self.rid_to_state.pop(recv_obj.rid, None)
+ state = self.rid_to_state[recv_obj.rid]
+ state.finished = True
+ state.out_list.append({
+ "text": "",
+ "meta_info": {
+ "id": recv_obj.rid,
+ "finish_reason": {
+ "type": "abort",
+ "message": "Abort before prefill",
+ },
+ "prompt_tokens": 1,
+ "completion_tokens": 0,
+ },
+ })
+ state.event.set()
def _handle_open_session_req_output(self, recv_obj):
self.session_futures[recv_obj.session_id].set_result(
diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py
index 786a34a1..4fe46338 100644
--- a/python/sglang/srt/managers/tp_worker.py
+++ b/python/sglang/srt/managers/tp_worker.py
@@ -38,6 +38,7 @@ from sglang.srt.managers.schedule_batch import ModelWorkerBatch, global_server_a
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool, TokenToKVPoolAllocator
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_executor.model_runner import ModelRunner
+from sglang.srt.patch_torch import monkey_patch_torch_reductions
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import MultiprocessingSerializer, broadcast_pyobj, set_random_seed
@@ -247,11 +248,12 @@ class TpModelWorker:
self, recv_req: UpdateWeightsFromDistributedReqInput
):
success, message = self.model_runner.update_weights_from_distributed(
- recv_req.name, recv_req.dtype, recv_req.shape
+ recv_req.names, recv_req.dtypes, recv_req.shapes, recv_req.group_name
)
return success, message
def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput):
+ monkey_patch_torch_reductions()
success, message = self.model_runner.update_weights_from_tensor(
named_tensors=MultiprocessingSerializer.deserialize(
recv_req.serialized_named_tensors[self.tp_rank]
diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py
index 995dedd0..b040f4b0 100644
--- a/python/sglang/srt/model_executor/model_runner.py
+++ b/python/sglang/srt/model_executor/model_runner.py
@@ -215,6 +215,7 @@ class ModelRunner:
self.support_pp = (
"pp_proxy_tensors" in inspect.signature(self.model.forward).parameters
)
+ self._model_update_group = {}
def initialize(self, min_per_gpu_memory: float):
server_args = self.server_args
@@ -704,7 +705,7 @@ class ModelRunner:
)
try:
- self._model_update_group = init_custom_process_group(
+ self._model_update_group[group_name] = init_custom_process_group(
backend=backend,
init_method=f"tcp://{master_address}:{master_port}",
world_size=world_size,
@@ -717,7 +718,7 @@ class ModelRunner:
logger.error(message)
return False, message
- def update_weights_from_distributed(self, name, dtype, shape):
+ def update_weights_from_distributed(self, names, dtypes, shapes, group_name):
"""
Update specific parameter in the model weights online
through `_model_update_group` process group.
@@ -727,19 +728,34 @@ class ModelRunner:
dtype: the data type of the parameter to be updated.
shape: the shape of the parameter to be updated.
"""
- target_dtype = (
- dtype if isinstance(dtype, torch.dtype) else getattr(torch, dtype)
- )
- assert (
- self._model_update_group is not None
- ), "model update group must be initialized"
+ assert group_name in self._model_update_group, (
+ f"Group {group_name} not in {list(self._model_update_group.keys)}. "
+ "Please call `init_weights_update_group` first."
+ )
try:
- weights = torch.empty(shape, dtype=target_dtype, device=self.device)
- torch.distributed.broadcast(weights, src=0, group=self._model_update_group)
- self.model.load_weights([(name, weights)])
- return True, f"Succeeded to update parameter {name} online."
+ weights = []
+ handles = []
+ for name, dtype, shape in zip(names, dtypes, shapes):
+ target_dtype = (
+ dtype if isinstance(dtype, torch.dtype) else getattr(torch, dtype)
+ )
+ weight = torch.empty(shape, dtype=target_dtype, device=self.device)
+ handles.append(
+ torch.distributed.broadcast(
+ weight,
+ src=0,
+ group=self._model_update_group[group_name],
+ async_op=True,
+ )
+ )
+ weights.append((name, weight))
+ for handle in handles:
+ handle.wait()
+
+ self.model.load_weights(weights)
+ return True, f"Succeeded to update parameter online."
except Exception as e:
error_msg = (
diff --git a/python/sglang/srt/openai_api/protocol.py b/python/sglang/srt/openai_api/protocol.py
index 351c1c56..180dd877 100644
--- a/python/sglang/srt/openai_api/protocol.py
+++ b/python/sglang/srt/openai_api/protocol.py
@@ -210,7 +210,7 @@ class CompletionResponseStreamChoice(BaseModel):
index: int
text: str
logprobs: Optional[LogProbs] = None
- finish_reason: Optional[Literal["stop", "length", "content_filter"]] = None
+ finish_reason: Optional[Literal["stop", "length", "content_filter", "abort"]] = None
matched_stop: Union[None, int, str] = None
@@ -444,7 +444,7 @@ class ChatCompletionResponseStreamChoice(BaseModel):
delta: DeltaMessage
logprobs: Optional[Union[LogProbs, ChoiceLogprobs]] = None
finish_reason: Optional[
- Literal["stop", "length", "tool_calls", "content_filter", "function_call"]
+ Literal["stop", "length", "tool_calls", "content_filter", "function_call", "abort"]
] = None
matched_stop: Union[None, int, str] = None
diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py
index ac04cdc7..d087745e 100644
--- a/python/sglang/srt/server_args.py
+++ b/python/sglang/srt/server_args.py
@@ -47,6 +47,7 @@ class ServerArgs:
tokenizer_path: Optional[str] = None
tokenizer_mode: str = "auto"
skip_tokenizer_init: bool = False
+ skip_warmup: bool = False
load_format: str = "auto"
trust_remote_code: bool = False
dtype: str = "auto"
@@ -122,6 +123,8 @@ class ServerArgs:
# Multi-node distributed serving
dist_init_addr: Optional[str] = None
+ nccl_port: Optional[int] = None
+ other_ports: Optional[List[int]] = None
nnodes: int = 1
node_rank: int = 0
@@ -538,6 +541,11 @@ class ServerArgs:
action="store_true",
help="If set, skip init tokenizer and pass input_ids in generate request.",
)
+ parser.add_argument(
+ "--skip-warmup",
+ action="store_true",
+ help="If set, skip warmup.",
+ )
parser.add_argument(
"--load-format",
type=str,
@@ -1492,6 +1500,9 @@ class ServerArgs:
help="Set multimodal attention backend.",
)
+ parser.add_argument("--nccl-port", type=int, default=None)
+ parser.add_argument("--other-ports", type=int, nargs="+", default=None)
+
@classmethod
def from_cli_args(cls, args: argparse.Namespace):
args.tp_size = args.tensor_parallel_size
@@ -1580,14 +1591,17 @@ class PortArgs:
@staticmethod
def init_new(server_args, dp_rank: Optional[int] = None) -> "PortArgs":
- port = server_args.port + random.randint(100, 1000)
- while True:
- if is_port_available(port):
- break
- if port < 60000:
- port += 42
- else:
- port -= 43
+ if server_args.nccl_port is None:
+ port = server_args.port + random.randint(100, 1000)
+ while True:
+ if is_port_available(port):
+ break
+ if port < 60000:
+ port += 42
+ else:
+ port -= 43
+ else:
+ port = server_args.nccl_port
if not server_args.enable_dp_attention:
# Normal case, use IPC within a single node
@@ -1613,6 +1627,27 @@ class PortArgs:
), "please provide --dist-init-addr as host:port of head node"
dist_init_host, dist_init_port = dist_init_addr
+
+ if server_args.other_ports is not None:
+ assert port not in server_args.other_ports
+ tokenizer_port = server_args.other_ports[0]
+ detokenizer_port = server_args.other_ports[1]
+ rpc_ipc_port = server_args.other_ports[2]
+ if dp_rank is None:
+ scheduler_input_port = server_args.other_ports[
+ 3
+ ] # TokenizerManager to DataParallelController
+ else:
+ scheduler_input_port = server_args.other_ports[3 + 1 + dp_rank]
+
+ return PortArgs(
+ tokenizer_ipc_name=f"tcp://{dist_init_host}:{tokenizer_port}",
+ scheduler_input_ipc_name=f"tcp://{dist_init_host}:{scheduler_input_port}",
+ detokenizer_ipc_name=f"tcp://{dist_init_host}:{detokenizer_port}",
+ nccl_port=port,
+ rpc_ipc_name=f"tcp://{dist_init_host}:{rpc_ipc_port}",
+ )
+
port_base = int(dist_init_port) + 1
if dp_rank is None:
scheduler_input_port = (
diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py
index bc0b50f3..4dbd1cfa 100644
--- a/python/sglang/srt/speculative/eagle_worker.py
+++ b/python/sglang/srt/speculative/eagle_worker.py
@@ -11,6 +11,10 @@ from sglang.srt.distributed import GroupCoordinator, patch_tensor_parallel_group
from sglang.srt.layers.dp_attention import disable_dp_size
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.sampler import get_token_ids_logprobs, get_top_logprobs
+from sglang.srt.managers.io_struct import (
+ UpdateWeightsFromDistributedReqInput,
+ UpdateWeightsFromTensorReqInput,
+)
from sglang.srt.managers.schedule_batch import (
ScheduleBatch,
get_last_loc,
@@ -22,6 +26,7 @@ from sglang.srt.model_executor.forward_batch_info import (
ForwardBatch,
ForwardMode,
)
+from sglang.srt.patch_torch import monkey_patch_torch_reductions
from sglang.srt.server_args import ServerArgs
from sglang.srt.speculative.build_eagle_tree import build_tree_kernel_efficient
from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
@@ -39,7 +44,13 @@ from sglang.srt.speculative.eagle_utils import (
select_top_k_tokens,
)
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
-from sglang.srt.utils import empty_context, fast_topk, get_available_gpu_memory, is_cuda
+from sglang.srt.utils import (
+ MultiprocessingSerializer,
+ empty_context,
+ fast_topk,
+ get_available_gpu_memory,
+ is_cuda,
+)
if is_cuda():
from sgl_kernel import segment_packbits
@@ -772,6 +783,29 @@ class EAGLEWorker(TpModelWorker):
logger.error("Detected errors during sampling! NaN in the logits.")
raise ValueError("Detected errors during sampling! NaN in the logits.")
+ def update_weights_from_distributed(
+ self, recv_req: UpdateWeightsFromDistributedReqInput
+ ):
+ raise NotImplementedError(
+ "EAGLEWorker does not support update_weights_from_distributed."
+ )
+
+ def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput):
+ monkey_patch_torch_reductions()
+ named_tensors = MultiprocessingSerializer.deserialize(
+ recv_req.serialized_named_tensors[self.tp_rank]
+ )
+ success, message = self.model_runner.update_weights_from_tensor(
+ named_tensors=named_tensors,
+ load_format=recv_req.load_format,
+ )
+ success, message = self.target_worker.model_runner.update_weights_from_tensor(
+ named_tensors=named_tensors,
+ load_format=recv_req.load_format,
+ )
+
+ return success, message
+
def load_token_map(token_map_path: str) -> List[int]:
if not os.path.exists(token_map_path):
diff --git a/python/sglang/srt/utils.py b/python/sglang/srt/utils.py
index a669b430..9ab99ba4 100644
--- a/python/sglang/srt/utils.py
+++ b/python/sglang/srt/utils.py
@@ -1787,8 +1787,16 @@ def get_ip() -> str:
def get_open_port() -> int:
port = os.getenv("SGLANG_PORT")
- if port is not None:
- port = int(port)
+ used_port = os.getenv("SGLANG_USED_PORT")
+ if used_port is not None:
+ used_port = list(map(int, used_port.split(",")))
+
+ if port is not None or used_port is not None:
+ if used_port is not None:
+ port = max(used_port) + random.randint(1, 1000)
+ else:
+ port = int(port)
+
while True:
try:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
diff --git a/sgl-router/py_src/sglang_router/launch_router.py b/sgl-router/py_src/sglang_router/launch_router.py
index 4f036a25..3e0cf9dd 100644
--- a/sgl-router/py_src/sglang_router/launch_router.py
+++ b/sgl-router/py_src/sglang_router/launch_router.py
@@ -43,6 +43,7 @@ class RouterArgs:
max_payload_size: int = 4 * 1024 * 1024 # 4MB
verbose: bool = False
log_dir: Optional[str] = None
+ log_level: Optional[str] = None
# Service discovery configuration
service_discovery: bool = False
selector: Dict[str, str] = dataclasses.field(default_factory=dict)
@@ -157,6 +158,13 @@ class RouterArgs:
default=None,
help="Directory to store log files. If not specified, logs are only output to console.",
)
+ parser.add_argument(
+ f"--{prefix}log-level",
+ type=str,
+ default=None,
+ choices=["debug", "info", "warning", "error", "critical"],
+ help="Set the logging level. If not specified, defaults to INFO.",
+ )
parser.add_argument(
f"--{prefix}service-discovery",
action="store_true",
@@ -225,6 +233,7 @@ class RouterArgs:
max_payload_size=getattr(args, f"{prefix}max_payload_size"),
verbose=getattr(args, f"{prefix}verbose", False),
log_dir=getattr(args, f"{prefix}log_dir", None),
+ log_level=getattr(args, f"{prefix}log_level", None),
service_discovery=getattr(args, f"{prefix}service_discovery", False),
selector=cls._parse_selector(getattr(args, f"{prefix}selector", None)),
service_discovery_port=getattr(args, f"{prefix}service_discovery_port"),
@@ -292,6 +301,7 @@ def launch_router(args: argparse.Namespace) -> Optional[Router]:
max_payload_size=router_args.max_payload_size,
verbose=router_args.verbose,
log_dir=router_args.log_dir,
+ log_level=router_args.log_level,
service_discovery=router_args.service_discovery,
selector=router_args.selector,
service_discovery_port=router_args.service_discovery_port,
diff --git a/sgl-router/py_src/sglang_router/router.py b/sgl-router/py_src/sglang_router/router.py
index c189cd58..08fde994 100644
--- a/sgl-router/py_src/sglang_router/router.py
+++ b/sgl-router/py_src/sglang_router/router.py
@@ -32,6 +32,7 @@ class Router:
max_tree_size: Maximum size of the approximation tree for cache-aware routing. Default: 2^24
verbose: Enable verbose logging. Default: False
log_dir: Directory to store log files. If None, logs are only output to console. Default: None
+ log_level: Logging level. Options: 'debug', 'info', 'warning', 'error', 'critical'.
service_discovery: Enable Kubernetes service discovery. When enabled, the router will
automatically discover worker pods based on the selector. Default: False
selector: Dictionary mapping of label keys to values for Kubernetes pod selection.
@@ -60,6 +61,7 @@ class Router:
max_payload_size: int = 4 * 1024 * 1024, # 4MB
verbose: bool = False,
log_dir: Optional[str] = None,
+ log_level: Optional[str] = None,
service_discovery: bool = False,
selector: Dict[str, str] = None,
service_discovery_port: int = 80,
@@ -85,6 +87,7 @@ class Router:
max_payload_size=max_payload_size,
verbose=verbose,
log_dir=log_dir,
+ log_level=log_level,
service_discovery=service_discovery,
selector=selector,
service_discovery_port=service_discovery_port,
diff --git a/sgl-router/py_test/test_launch_router.py b/sgl-router/py_test/test_launch_router.py
index c6f0444f..990f37a0 100644
--- a/sgl-router/py_test/test_launch_router.py
+++ b/sgl-router/py_test/test_launch_router.py
@@ -38,6 +38,7 @@ class TestLaunchRouter(unittest.TestCase):
max_payload_size=4 * 1024 * 1024, # 4MB
verbose=False,
log_dir=None,
+ log_level=None,
service_discovery=False,
selector=None,
service_discovery_port=80,
diff --git a/sgl-router/src/lib.rs b/sgl-router/src/lib.rs
index 4915d3c5..71a411a1 100644
--- a/sgl-router/src/lib.rs
+++ b/sgl-router/src/lib.rs
@@ -33,6 +33,7 @@ struct Router {
max_payload_size: usize,
verbose: bool,
log_dir: Option<String>,
+ log_level: Option<String>,
service_discovery: bool,
selector: HashMap<String, String>,
service_discovery_port: u16,
@@ -59,6 +60,7 @@ impl Router {
max_payload_size = 4 * 1024 * 1024,
verbose = false,
log_dir = None,
+ log_level = None,
service_discovery = false,
selector = HashMap::new(),
service_discovery_port = 80,
@@ -81,6 +83,7 @@ impl Router {
max_payload_size: usize,
verbose: bool,
log_dir: Option<String>,
+ log_level: Option<String>,
service_discovery: bool,
selector: HashMap<String, String>,
service_discovery_port: u16,
@@ -103,6 +106,7 @@ impl Router {
max_payload_size,
verbose,
log_dir,
+ log_level,
service_discovery,
selector,
service_discovery_port,
@@ -164,6 +168,7 @@ impl Router {
verbose: self.verbose,
max_payload_size: self.max_payload_size,
log_dir: self.log_dir.clone(),
+ log_level: self.log_level.clone(),
service_discovery_config,
prometheus_config,
})
diff --git a/sgl-router/src/server.rs b/sgl-router/src/server.rs
index 0d6cf691..7ebdc442 100644
--- a/sgl-router/src/server.rs
+++ b/sgl-router/src/server.rs
@@ -161,6 +161,7 @@ pub struct ServerConfig {
pub verbose: bool,
pub max_payload_size: usize,
pub log_dir: Option<String>,
+ pub log_level: Option<String>,
pub service_discovery_config: Option<ServiceDiscoveryConfig>,
pub prometheus_config: Option<PrometheusConfig>,
}
@@ -174,7 +175,21 @@ pub async fn startup(config: ServerConfig) -> std::io::Result<()> {
level: if config.verbose {
Level::DEBUG
} else {
- Level::INFO
+ config.log_level
+ .as_deref()
+ .and_then(|s| {
+ match s.to_uppercase().parse::<Level>() {
+ Ok(l) => Some(l),
+ Err(_) => {
+ warn!(
+ "Invalid log level string: '{}'. Defaulting to INFO (as verbose is false).",
+ s
+ );
+ None
+ }
+ }
+ })
+ .unwrap_or(Level::INFO)
},
json_format: false,
log_dir: config.log_dir.clone(),
docs/en/amd_toturial.md
+292
View File
@@ -0,0 +1,292 @@
# AMD Toturial
## Introduction
If you are running Slime on AMD's Instinct, please refer to the following materials. This tutorial will explain how to set up the development environment (Docker), use the modified AMD dependencies, and provide an example for running the experiments. The current rocm docker only support AMD's MI300.
⚠️ Note: Feel free to reach out [Yusheng Su](https://yushengsu-thu.github.io/), if you encounter problems on AMD instinct.
<!-- First, you need to configure the Slime runtime environment according to the [Readme](../../README.md) documentation and cd to the Slime project directory. -->
## Docker
You can download the prebuilt image from DockerHub: [yushengsuthu/slime-amd]().
```bash
docker pull yushengsuthu/slime-amd:slime_ubuntu22.04_rocm6.3.4-patch-numa_vllm0.8.5-patch_sglang0.4.7_megatron-core-patch_ray0.47-patch
```
Or you can use the [Dockerfile.rocm](docker/Dockerfile.rocm) to build it on your side.
```bash
cd docker
docker build -f Dockerfile.rocm -t slime_ubuntu22.04_rocm6.3.4-patch-numa_vllm0.8.5-patch_sglang0.4.7_megatron-core-patch_ray0.47-patch
```
## Quick Start
### Environment Setup
Based on the [yushengsuthu/slime-amd]() image (pre-installed with SGLang 0.4.7 and Megatron):
```bash
docker run --rm -it \
--device /dev/dri \
--device /dev/kfd \
-p 8265:8265 \
--group-add video \
--cap-add SYS_PTRACE \
--security-opt seccomp=unconfined \
--privileged \
-v $HOME/.ssh:/root/.ssh \
-v $HOME:$HOME \
--shm-size 128G \
--name slime_dev \
--ulimit memlock=-1 \
--ulimit stack=67108864 \
-w $PWD \
slime_ubuntu22.04_rocm6.3.4-patch-numa_vllm0.8.5-patch_sglang0.4.7_megatron-core-patch_ray0.47-patch:latest \
/bin/bash
```
Then, download and install Slime.
```bash
git clone https://github.com/THUDM/slime.git
cd slime
pip install -e .
```
### Checkpoint Format Conversion
Since slime uses Megatron, and Megatron does not support loading Hugging Face checkpoints directly, we need to convert the model to the `torch_dist` format that Megatron supports.
#### HF → Megatron torch\_dist ckpt
Use [mbridge](https://github.com/ISEEKYAN/mbridge.git) or [Megatron-LM-amd_version-amd](https://github.com/yushengsu-thu/Megatron-LM-amd_version.git) for conversion:
```bash
cd slime/
# PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \
MEGATRON_LM_PATH=$(pip list | grep megatron-core | awk '{print $NF}')
PYTHONPATH=${MEGATRON_LM_PATH} python tools/convert_hf_to_torch_dist.py \
--hf-checkpoint model/Qwen3-4B \
--save model/Qwen3-4B_torch_dist
```
⚠️ If you encounter an issue where slime cannot be found, please run `pip install -e .` in the slime directory.
### Example: Qwen3-4B
We provide examples to use [Qwen3-4B](https://huggingface.co/Qwen/Qwen3-4B), please refer to:
- [Example: Qwen3-4B Model](scripts/run-qwen3-4B-amd.sh): Just run `scripts/run-qwen3-4B-amd.sh`
We show the training script below.
⚠️ TODO: The [AMD-version torch_memory_saver](https://github.com/YangWang92/torch_memory_saver_numa/tree/numa) does not seem to clear memory properly; thus, we set `--sglang-mem-fraction-static` as `0.4` currently. We will continue investigating and focus on ROCm's virtual memory management for further modifications.
⚠️ TODO: ROCM seems to not support `apex` yet. Thus, we need to disable `--no-gradient-accumulation-fusion` currently. We will continue investigating how to enable this.
⚠️ Note: The main difference between AMD's training script and NVIDIA's script is that you need to set `RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES` and `HIP_VISIBLE_DEVICES` for ray to function properly on AMD GPUs.
```bash
#!/bin/bash
####clear before training
pkill -9 sglang
sleep 3
ray stop --force
pkill -9 ray
pkill -9 python
sleep 3
pkill -9 ray
pkill -9 python
set -euxo pipefail
### AMD Support ###
SLIME_DIR="/home/yushensu/projects/slime" # Need to change to your own path
export SLIME_DIR=$SLIME_DIR
MODEL_DIR="/home/yushensu/projects/model" # Need to change to your own path
export MODEL_DIR=$MODEL_DIR
DATA_DIR="/home/yushensu/projects/data" # Need to change to your own path
export DATA_DIR=$DATA_DIR
# For AMD GPU
export RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=${RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES:-"1"} # Must set to 1
export HIP_VISIBLE_DEVICES=${HIP_VISIBLE_DEVICES:-"0,1,2,3,4,5,6,7"} #You can choose which gpus to use
####################
# ### AMD Support ### (If you do not istall, please install them)
# # # Clone and install Megatron-LMi-amd_version
# export MAX_JOBS=512
# cd $SLIME_DIR
# pip uninstall megatron-core -y
# if [ ! -d "Megatron-LM-amd_version" ]; then
# git clone git@github.com:yushengsu-thu/Megatron-LM-amd_version.git
# else
# echo "Megatron-LM-amd_version directory already exists, skipping clone"
# fi
# cd Megatron-LM-amd_version
# pip install -vvv -e .
# cd $SLIME_DIR
# # Install slime
# pip install -e .
# ####################
# will prevent ray from buffering stdout/stderr
export PYTHONBUFFERED=16
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
source "${SCRIPT_DIR}/models/qwen3-4B.sh"
CKPT_ARGS=(
--hf-checkpoint ${MODEL_DIR}/Qwen3-4B
#--hf-checkpoint /root/Qwen3-4B-FP8
--ref-load ${MODEL_DIR}/Qwen3-4B_torch
--load ${MODEL_DIR}/Qwen3-4B_slime/
--save ${MODEL_DIR}/Qwen3-4B_slime/
--save-interval 20
)
ROLLOUT_ARGS=(
--prompt-data ${DATA_DIR}/dapo-math-17k/dapo-math-17k.jsonl
--input-key prompt
--label-key label
--apply-chat-template
--rollout-shuffle
--rm-type deepscaler
--num-rollout 3000
--rollout-batch-size 32
--n-samples-per-prompt 8
--rollout-max-response-len 8192
--rollout-temperature 0.8
--global-batch-size 256
--balance-data
)
EVAL_ARGS=(
--eval-interval 20
--eval-prompt-data aime ${DATA_DIR}/aime-2024/aime-2024.jsonl
--n-samples-per-eval-prompt 16
--eval-max-response-len 16384
--eval-top-p 0.7
)
PERF_ARGS=(
--tensor-model-parallel-size 1
--sequence-parallel
--pipeline-model-parallel-size 1
--context-parallel-size 1
--expert-model-parallel-size 1
--expert-tensor-parallel-size 1
--recompute-granularity full
--recompute-method uniform
--recompute-num-layers 1
# --micro-batch-size 1
--use-dynamic-batch-size
--max-tokens-per-gpu 9216
)
GRPO_ARGS=(
--advantage-estimator grpo
--use-kl-loss
--kl-loss-coef 0.00
--kl-loss-type low_var_kl
--kl-coef 0.00
--entropy-coef 0.00
--eps-clip 0.2
--eps-clip-high 0.28
)
OPTIMIZER_ARGS=(
--optimizer adam
--lr 1e-6
--lr-decay-style constant
--weight-decay 0.1
--adam-beta1 0.9
--adam-beta2 0.98
)
WANDB_ARGS=(
#--use-wandb
# --wandb-project slime-dev
# --wandb-group qwen3-4B-test
# --wandb-key ${WANDB_KEY}
)
### AMD Support ###
# Need to fix some issue with torch_memory_saver in rocm to support larger --sglang-mem-fraction-static
# SGLANG_ARGS=(
# --rollout-num-gpus-per-engine 2
# --sglang-mem-fraction-static 0.7
# )
SGLANG_ARGS=(
--rollout-num-gpus-per-engine 2
--sglang-mem-fraction-static 0.4
)
####################
MISC_ARGS=(
# default dropout in megatron is 0.1
--attention-dropout 0.0
--hidden-dropout 0.0
# should be good for model performance
--accumulate-allreduce-grads-in-fp32
--attention-softmax-in-fp32
# need to comment this when using model with MLA
--attention-backend flash
### AMD Support ###
# disable gradient accumulation fusion: Need to add apex to enable this
--no-gradient-accumulation-fusion
###################
)
# launch the master node of ray in container
export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
export MASTER_PORT=${MASTER_PORT:-"12345"}
NUM_GPUS=$(echo ${HIP_VISIBLE_DEVICES} | tr ',' '\n' | wc -l)
ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus ${NUM_GPUS} --disable-usage-stats
# "PYTHONPATH": "$(dirname $(python3 -c 'import megatron.core; print(megatron.core.__file__)'))"
ray job submit --address="http://127.0.0.1:8265" \
--runtime-env-json='{
"env_vars": {
"PYTHONPATH": "'${SLIME_DIR}'/Megatron-LM-amd_version/",
"CUDA_DEVICE_MAX_CONNECTIONS": "1"
}
}' \
-- python3 train.py \
--actor-num-nodes 1 \
--actor-num-gpus-per-node 8 \
--colocate \
${MODEL_ARGS[@]} \
${CKPT_ARGS[@]} \
${ROLLOUT_ARGS[@]} \
${OPTIMIZER_ARGS[@]} \
${GRPO_ARGS[@]} \
${DISTRIBUTED_ARGS[@]} \
${WANDB_ARGS[@]} \
${PERF_ARGS[@]} \
${EVAL_ARGS[@]} \
${SGLANG_ARGS[@]} \
${MISC_ARGS[@]}
####clear after training
pkill -9 sglang
sleep 3
ray stop --force
pkill -9 ray
pkill -9 python
sleep 3
pkill -9 ray
pkill -9 python
```
scripts/run-qwen3-4B-amd.sh
+21 -17
View File
@@ -34,23 +34,23 @@ export HIP_VISIBLE_DEVICES=${HIP_VISIBLE_DEVICES:-"0,1,2,3,4,5,6,7"} #You can ch
####################
### AMD Support ### (If you do not istall, please install them)
# # Clone and install Megatron-LMi-amd_version
export MAX_JOBS=512
cd $SLIME_DIR
pip uninstall megatron-core -y
if [ ! -d "Megatron-LM-amd_version" ]; then
git clone git@github.com:yushengsu-thu/Megatron-LM-amd_version.git
else
echo "Megatron-LM-amd_version directory already exists, skipping clone"
fi
cd Megatron-LM-amd_version
pip install -vvv -e .
cd $SLIME_DIR
# ### AMD Support ### (If you do not istall, please install them)
# # # Clone and install Megatron-LMi-amd_version
# export MAX_JOBS=512
# cd $SLIME_DIR
# pip uninstall megatron-core -y
# if [ ! -d "Megatron-LM-amd_version" ]; then
# git clone git@github.com:yushengsu-thu/Megatron-LM-amd_version.git
# else
# echo "Megatron-LM-amd_version directory already exists, skipping clone"
# fi
# cd Megatron-LM-amd_version
# pip install -vvv -e .
# cd $SLIME_DIR
# Install slime
pip install -e .
####################
# # Install slime
# pip install -e .
# ####################
@@ -176,10 +176,14 @@ export MASTER_PORT=${MASTER_PORT:-"12345"}
NUM_GPUS=$(echo ${HIP_VISIBLE_DEVICES} | tr ',' '\n' | wc -l)
ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus ${NUM_GPUS} --disable-usage-stats
# "PYTHONPATH": "/workspace/Megatron-LM-amd_version/",
MEGATRON_LM_PATH=$(pip list | grep megatron-core | awk '{print $NF}')
ray job submit --address="http://127.0.0.1:8265" \
--runtime-env-json='{
"env_vars": {
"PYTHONPATH": "'${SLIME_DIR}'/Megatron-LM-amd_version/",
"PYTHONPATH": "'$MEGATRON_LM_PATH'",
"CUDA_DEVICE_MAX_CONNECTIONS": "1"
}
}' \
slime/ray/ppo_actor.py
+5 -2
View File
@@ -7,9 +7,12 @@ from typing import Dict
import ray
import torch
import torch.distributed as dist
# from cumem_allocator import CuMemAllocator
### AMD Support ###
from vllm.device_allocator.cumem import CuMemAllocator
if torch.version.hip: # Check if using AMD/ROCm
from vllm.device_allocator.cumem import CuMemAllocator
else: # NVIDIA/CUDA
from cumem_allocator import CuMemAllocator
###################