Fleet/all

[dzorlu]

All

---

[gemini-code-assist]

Code Review

This pull request implements Fleet-specific optimizations for multi-node FSDP2 training, featuring memory-efficient chunked forward passes, synchronous model offloading with barriers, and support for multimodal Vision-Language tasks. It introduces a new task generation environment with database exploration tools and provides comprehensive S3 checkpoint management. The review feedback highlights a documentation inconsistency regarding flash attention usage, identifies deprecated Pydantic configuration patterns, warns against blocking calls within asynchronous functions, and suggests refactoring duplicated tool parsing logic into a shared utility.

[gemini-code-assist]

There appears to be a documentation inconsistency. This line states that flash_attn=true is part of the working configuration at 72K sequence length. However, integrations/fleet/CHANGELOG.md (line 60) and scripts/fleet-35b-run.sh (line 43) specify that flash_attn=false is used. Please update this documentation for consistency.

[gemini-code-assist]

The inner Config class for Pydantic model configuration is deprecated in Pydantic v2 and will be removed in v3. Configuration should be passed as keyword arguments directly to the BaseModel class definition. You should remove this Config class and change line 102 to class RolloutOutput(BaseModel, frozen=True):.

[gemini-code-assist]

The use of .result() here is a blocking call inside an async function, which can hinder performance by blocking the event loop. If the tinker SDK provides an asynchronous way to get the result (e.g., an await-able method), it should be used instead to maintain the benefits of an async architecture. This also applies to the .result() calls on lines 877-878.

[gemini-code-assist]

This file is nearly identical to skyrl-gym/skyrl_gym/envs/fleet_task/tool_call_parser.py. Duplicating this utility makes the code harder to maintain, as bug fixes or improvements would need to be applied in multiple places. Consider moving this parsing logic to a shared utility module within skyrl_gym to avoid duplication.

[devin-ai-integration]

✅ Devin Review: No Issues Found

Devin Review analyzed this PR and found no bugs or issues to report.

[devin-ai-integration]

Devin Review found 2 new potential issues.

View 7 additional findings in Devin Review.

[devin-ai-integration]

🔴 loss_chunk_size=None from config causes TypeError in model forward pass

The config default for loss_chunk_size is Optional[int] = None (config.py:623) and the YAML default is null (ppo_base_config.yaml:260). fsdp_worker.py:191 passes self.cfg.loss_chunk_size (which is None) directly to HFModelWrapper(loss_chunk_size=None). Since None is passed explicitly, it overrides the wrapper's default of 0, setting self.loss_chunk_size = None. Then at line 388, use_chunked = self.loss_chunk_size > 0 raises TypeError: '>' not supported between instances of 'NoneType' and 'int', crashing the first training forward pass for any run that doesn't explicitly set loss_chunk_size to an integer (i.e., all non-Fleet-35B runs using default config).

Suggested change

	use_chunked = self.loss_chunk_size > 0
	use_chunked = self.loss_chunk_size is not None and self.loss_chunk_size > 0

---

Was this helpful? React with 👍 or 👎 to provide feedback.

[devin-ai-integration]

🟡 stage_chunks fallback to dp_size silently drops trailing samples

When the while loop at line 200 exhausts without finding a valid mini_batch_size (because no multiple of dp_size evenly divides len(data)), it falls back to mini_batch_size = dp_size at line 203. If len(data) % dp_size != 0, the subsequent num_mini_batches = len(data) // mini_batch_size at line 212 silently drops the remainder samples. For example, with len(data)=145 and dp_size=16, 145 // 16 = 9 mini-batches process only 144 samples, silently dropping 1. This can occur when hint augmentation produces batch sizes that aren't multiples of dp_size.

Suggested change

	if mini_batch_size <= 0:
	mini_batch_size = dp_size
	if mini_batch_size <= 0:
	mini_batch_size = dp_size
	if len(data) % mini_batch_size != 0:
	logger.warning(
	f"Fallback mini_batch_size={mini_batch_size} does not evenly divide "
	f"batch of {len(data)} samples; {len(data) % mini_batch_size} samples will be dropped."
	)

---

Was this helpful? React with 👍 or 👎 to provide feedback.

[devin-ai-integration]

Devin Review found 1 new potential issue.

View 11 additional findings in Devin Review.

[devin-ai-integration]

🟡 Hint augmentation base_rep_id uses batch index instead of repetition_id, creating collisions across groups

base_rep_id = max(item[0] for item in items) + 1 at line 866 uses item[0], which is the index i into all_outputs (the batch position, e.g. 0–127), not the repetition_id from TrajectoryID. Since prepare_generator_input (skyrl/train/generators/utils.py:404) creates repetition_id values 0 through n_samples_per_prompt-1, hinted TrajectoryID objects should use repetition_id >= n_samples_per_prompt. Instead, base_rep_id is a large batch index (e.g., 120 for a prompt at batch position 120), which is semantically wrong and could collide with batch indices from other instance groups.

Fix

Replace max(item[0] for item in items) with max(trajectory_ids[item[0]].repetition_id for item in items) to correctly derive from the actual repetition_id.

Suggested change

	base_rep_id = max(item[0] for item in items) + 1
	base_rep_id = max(trajectory_ids[item[0]].repetition_id for item in items) + 1

---

Was this helpful? React with 👍 or 👎 to provide feedback.

[devin-ai-integration]

Devin Review found 2 new potential issues.

View 14 additional findings in Devin Review.

[devin-ai-integration]

🔴 Chunked lm_head identity replacement silently fails with LoRA-wrapped models

When loss_chunk_size > 0, _IdentityLMHead is assigned via self.model.lm_head = _IdentityLMHead() (model_wrapper.py:394). For non-LoRA models, self.model is the CausalLM directly, so the identity replacement works. But when lora_rank > 0, self.model is a PeftModel. Python's __setattr__ sets lm_head on the PeftModel instance dict, but PeftModel's forward calls self.base_model(...) → inner CausalLM, which still accesses its original lm_head (not the identity). The chunked path silently falls back to materializing full (B, S, vocab_size) logits, defeating the memory optimization. This currently doesn't trigger since fleet configs use lora_rank=0, but it's a correctness trap for anyone enabling LoRA with loss_chunk_size.

How the bug manifests

With LoRA enabled:

1. self.model.lm_head (getter) → PeftModel __getattr__ proxies to inner model's lm_head ✓
2. self.model.lm_head = _IdentityLMHead() (setter) → sets on PeftModel instance dict, NOT inner model
3. Forward: PeftModel → base_model → inner CausalLM uses original lm_head → full logits materialized
4. output["logits"] has shape (B, S, vocab_size) instead of (B, S, hidden_dim)
5. _chunked_lm_head_forward receives logits instead of hidden states → garbage logprobs

Prompt for agents

In skyrl/backends/skyrl_train/workers/model_wrapper.py, the _IdentityLMHead replacement at lines 393-394 does not work when self.model is a PeftModel (LoRA enabled). The fix should unwrap through PeftModel to set the identity on the actual inner model. For example:

1. Add a helper to find the actual CausalLM model inside PeftModel wrappers:
   inner_model = self.model
   while hasattr(inner_model, 'base_model') or hasattr(inner_model, 'model'):
       inner_model = getattr(inner_model, 'model', inner_model)

2. Replace lm_head on the inner_model:
   inner_model.lm_head = _IdentityLMHead()

3. Restore on the same inner_model in the finally block.

Alternatively, add an assertion at lines 388-390 that raises an error if both loss_chunk_size > 0 and the model is a PeftModel, to prevent silent incorrect behavior.

---

Was this helpful? React with 👍 or 👎 to provide feedback.

[devin-ai-integration]

🟡 export_tasks.py crashes with FileNotFoundError when output path has no directory component

os.makedirs(os.path.dirname(os.path.expanduser(output_file)), exist_ok=True) at export_tasks.py:51 will raise FileNotFoundError if output_file is a bare filename like "tasks.json". os.path.dirname("tasks.json") returns "", and os.makedirs("", exist_ok=True) fails. The fix on line 53 (output_path = os.path.expanduser(output_file)) handles ~ correctly but doesn't prevent the earlier makedirs crash.

Suggested change

	os.makedirs(os.path.dirname(os.path.expanduser(output_file)), exist_ok=True)
	output_path = os.path.expanduser(output_file)
	output_dir = os.path.dirname(output_path)
	if output_dir:
	os.makedirs(output_dir, exist_ok=True)

---

Was this helpful? React with 👍 or 👎 to provide feedback.