73 lines
2.9 KiB
Bash
Executable File
73 lines
2.9 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
set -euo pipefail
|
|
|
|
MEGATRON_PATH="${MEGATRON_PATH:-/opt/Megatron-Bridge/3rdparty/Megatron-LM}"
|
|
SERVER_PY="${MEGATRON_PATH}/tools/run_text_generation_server.py"
|
|
ENGINE_PY="${MEGATRON_PATH}/megatron/core/inference/text_generation_server/run_mcore_engine.py"
|
|
|
|
python3 - <<'PY'
|
|
import importlib.util
|
|
missing = [name for name in ["flask", "flask_restful"] if importlib.util.find_spec(name) is None]
|
|
if missing:
|
|
raise SystemExit(f"missing python packages: {missing}; install with: pip install flask-restful")
|
|
print("flask dependencies ok")
|
|
PY
|
|
|
|
python3 - "$SERVER_PY" "$ENGINE_PY" <<'PY'
|
|
from pathlib import Path
|
|
import sys
|
|
|
|
server = Path(sys.argv[1])
|
|
engine = Path(sys.argv[2])
|
|
|
|
text = server.read_text()
|
|
old = "inference_context = StaticInferenceContext(args.inference_max_requests, args.inference_max_sequence_length)"
|
|
new = "inference_context = StaticInferenceContext(args.inference_max_requests, getattr(args, 'inference_max_sequence_length', getattr(args, 'inference_max_seq_length', args.seq_length)))"
|
|
if old in text:
|
|
text = text.replace(old, new)
|
|
server.write_text(text)
|
|
print(f"patched {server}: inference_max_sequence_length fallback")
|
|
elif new in text:
|
|
print(f"already patched {server}: inference_max_sequence_length fallback")
|
|
else:
|
|
print(f"warning: expected inference_context line not found in {server}")
|
|
|
|
text = engine.read_text()
|
|
old = '"tokens": [x.prompt_tokens + x.generated_tokens.tolist() for x in result],'
|
|
new = '"tokens": [x.prompt_tokens.tolist() + x.generated_tokens if hasattr(x.prompt_tokens, "tolist") else x.prompt_tokens + x.generated_tokens for x in result],'
|
|
if old in text:
|
|
text = text.replace(old, new)
|
|
engine.write_text(text)
|
|
print(f"patched {engine}: token list concatenation")
|
|
elif "x.prompt_tokens.tolist() + x.generated_tokens" in text:
|
|
print(f"already patched {engine}: token list concatenation")
|
|
else:
|
|
print(f"warning: expected token concat line not found in {engine}")
|
|
|
|
text = engine.read_text()
|
|
old = " skip_prompt_log_probs=False,"
|
|
new = " skip_prompt_log_probs=True,"
|
|
if old in text:
|
|
text = text.replace(old, new)
|
|
engine.write_text(text)
|
|
print(f"patched {engine}: skip prompt logprobs")
|
|
elif new in text:
|
|
print(f"already patched {engine}: skip prompt logprobs")
|
|
else:
|
|
print(f"warning: expected skip_prompt_log_probs line not found in {engine}")
|
|
|
|
text = engine.read_text()
|
|
old = " response_logprobs = [x.prompt_log_probs + x.generated_log_probs for x in result]"
|
|
new = " response_logprobs = [x.generated_log_probs for x in result]"
|
|
if old in text:
|
|
text = text.replace(old, new)
|
|
engine.write_text(text)
|
|
print(f"patched {engine}: generated-only logprobs response")
|
|
elif new in text:
|
|
print(f"already patched {engine}: generated-only logprobs response")
|
|
else:
|
|
print(f"warning: expected logprobs response line not found in {engine}")
|
|
PY
|
|
|
|
echo "Megatron inference patch setup complete."
|