Skip to content

Add support for running MaxText models (Qwen3.5) in DeepSWE training. #1827

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
Aug 5, 2026
Merged

Add support for running MaxText models (Qwen3.5) in DeepSWE training. #1827

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
96 changes: 73 additions & 23 deletions examples/deepswe/train_deepswe_nb.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -257,11 +257,25 @@
)

args, _ = parser.parse_known_args()

# Register MaxText vLLM adapter if using a MaxText model
if args.model_source == "maxtext":
try:
from maxtext.integration.vllm import maxtext_vllm_adapter # pylint: disable=g-import-not-at-top # pytype: disable=import-error
maxtext_vllm_adapter.register()
logging.info("Successfully registered MaxTextForCausalLM model with vLLM.")
except ImportError as e:
logging.warning("Could not import maxtext_vllm_adapter: %s", e)

MODEL_VERSION = args.model_version
NODE_SELECTOR_VAL = args.node_selector_val


# Monkeypatch r2egym DockerRuntime to dynamically configure Kubernetes nodeSelector.
# This is required because r2egym hardcodes the CPU nodepool name (using
# Karpenter bigcpu-standby), which does not exist in our GKE cluster. We
# override it here to match the nodepool configured via the
# --node_selector_val flag.
def patch_kubernetes_runtime():
try:
from r2egym.agenthub.runtime.docker import DockerRuntime
Expand Down Expand Up @@ -369,18 +383,8 @@ def patched_create_namespaced_pod(*args, **kwargs):
from tunix import PerfMetricsConfig
from tunix.perf.experimental.export import PerfMetricsExport
from tunix.rl.agentic.rewards.reward_types import RewardOutput
from examples.deepswe.swe_agent import (
SWE_SYSTEM_PROMPT,
SWE_SYSTEM_PROMPT_FN_CALL,
SWE_USER_PROMPT,
SWE_USER_PROMPT_FN_CALL,
SWEAGENT_SYSTEM_PROMPT,
SWEAGENT_USER_PROMPT,
)

# Assumed custom imports based on usage
from examples.deepswe.swe_agent import SWEAgent
from examples.deepswe.swe_env import SWEEnv
from examples.deepswe import swe_agent
from examples.deepswe import swe_env

# %%
# ==========================================
Expand Down Expand Up @@ -432,10 +436,9 @@ def patched_create_namespaced_pod(*args, **kwargs):
print(f"Model not found locally. Starting download to {MODEL_PATH}...")
os.makedirs(MODEL_PATH, exist_ok=True)

# Assumes "Qwen/" organization prefix for HF download. Adjust if using
# other models.
# Requires full HF repository ID (e.g. "Qwen/Qwen3-32B").
snapshot_download( # pyrefly: ignore[no-matching-overload]
repo_id=f"Qwen/{MODEL_VERSION}",
repo_id=MODEL_VERSION,
local_dir=MODEL_PATH,
local_dir_use_symlinks=False,
)
Expand Down Expand Up @@ -562,7 +565,7 @@ def patched_create_namespaced_pod(*args, **kwargs):
tokenizer_path = MODEL_PATH
local_files_only = True
if MODEL_SOURCE == "maxtext" and MODEL_PATH.startswith("gs://"):
tokenizer_path = f"Qwen/{MODEL_VERSION}"
tokenizer_path = MODEL_VERSION
local_files_only = False
print(f"Loading tokenizer from HF Hub: {tokenizer_path}")

Expand Down Expand Up @@ -769,6 +772,8 @@ def mixed_type_batch_fn(elements):
model_source=ModelSource.MAXTEXT,
model_path=MODEL_PATH,
enable_checkpointing=True,
allow_split_physical_axes=True,
scan_layers=False,
)
else:
qwen_reference = params_lib.create_model_from_safe_tensors(
Expand Down Expand Up @@ -808,6 +813,7 @@ def get_lora_model(base_model, model_mesh):
graph_def,
jax.tree.map(jnp.copy, params),
)

sft_utils.show_hbm_usage()

# %%
Expand Down Expand Up @@ -887,6 +893,28 @@ def get_lora_model(base_model, model_mesh):
},
}

if MODEL_SOURCE == "maxtext":
vllm_rollout_dict["rollout_vllm_kwargs"]["hf_overrides"] = {
"architectures": ["MaxTextForCausalLM"]
}
vllm_rollout_dict["rollout_vllm_additional_config"] = {
"maxtext_config": {
"model_name": MODEL_VERSION.lower().split("/")[-1],
"model_call_mode": "inference",
"enable_dp_attention": False,
"allow_split_physical_axes": True,
"log_config": False,
"weight_dtype": "bfloat16",
"prefuse_moe_weights": True,
}
}
# Force no-op mappings for weight sync if both trainer and sampler use MaxText
if hasattr(qwen_reference, "use_no_op_mappings"):
qwen_reference.use_no_op_mappings = True
if hasattr(qwen_actor, "use_no_op_mappings"):
qwen_actor.use_no_op_mappings = True
logging.info("Forced use_no_op_mappings=True on actor/reference models.")


if ROLLOUT_ENGINE == "sglang_jax":
rollout_engine_config = base_rollout.RolloutConfig(
Expand All @@ -907,6 +935,23 @@ def get_lora_model(base_model, model_mesh):
else:
raise ValueError(f"Unsupported rollout engine: {ROLLOUT_ENGINE}")


def filter_logical_rules(rules, mesh):
"""Filters logical sharding rules to keep only physical axes present in mesh."""
valid_axes = set(mesh.shape.keys())
filtered_rules = []
for logical_axis, physical_axes in rules:
if isinstance(physical_axes, (list, tuple)):
new_phys = [ax for ax in physical_axes if ax in valid_axes]
filtered_rules.append((logical_axis, tuple(new_phys)))
else:
if physical_axes in valid_axes:
filtered_rules.append((logical_axis, physical_axes))
else:
filtered_rules.append((logical_axis, ()))
return tuple(filtered_rules)


role_to_logical_axis_rule = None
logical_rules = getattr(
getattr(getattr(qwen_reference, "base", None), "config", None),
Expand All @@ -916,9 +961,15 @@ def get_lora_model(base_model, model_mesh):
if logical_rules:
print(f"Configuring role_to_logical_axis_rule with: {logical_rules}")
role_to_logical_axis_rule = {
rl_engine_lib.Role.ACTOR: logical_rules,
rl_engine_lib.Role.REFERENCE: logical_rules,
rl_engine_lib.Role.ROLLOUT: logical_rules,
rl_engine_lib.Role.ACTOR: filter_logical_rules(
logical_rules, train_mesh
),
rl_engine_lib.Role.REFERENCE: filter_logical_rules(
logical_rules, train_mesh
),
rl_engine_lib.Role.ROLLOUT: filter_logical_rules(
logical_rules, rollout_mesh
),
}

cluster_config = rl_engine_lib.ClusterConfig(
Expand Down Expand Up @@ -954,7 +1005,6 @@ def get_lora_model(base_model, model_mesh):
cluster_config=cluster_config,
)


# %%
# ==========================================
# 10. Learner & Agent Setup
Expand All @@ -966,7 +1016,7 @@ def get_lora_model(base_model, model_mesh):
"max_response_length": MAX_RESPONSE_LENGTH,
"beta": BETA,
"epsilon": EPSILON,
"system_prompt": SWE_SYSTEM_PROMPT,
"system_prompt": swe_agent.SWE_SYSTEM_PROMPT,
"max_concurrency": MAX_CONCURRENCY,
"epsilon_high": EPSILON_HIGH,
"off_policy_steps": OFF_POLICY_STEPS,
Expand All @@ -984,9 +1034,9 @@ def get_lora_model(base_model, model_mesh):
agentic_grpo_learner = agentic_grpo_learner.GRPOLearner(
rl_engine=rl_engine,
reward_fns=None,
agent_class=SWEAgent,
agent_class=swe_agent.SWEAgent,
agent_kwargs={},
env_class=SWEEnv,
env_class=swe_env.SWEEnv,
env_kwargs={
"max_steps": MAX_TURNS,
"step_timeout": STEP_TIMEOUT_SECS,
Expand Down
Loading