Skip to content

Optimize the implementation of Sampler and ProduceStrategy#1493

Merged
YanhuiDua merged 3 commits intoInternLM:rl_designfrom
YanhuiDua:optim_produce
Feb 27, 2026
Merged

Optimize the implementation of Sampler and ProduceStrategy#1493
YanhuiDua merged 3 commits intoInternLM:rl_designfrom
YanhuiDua:optim_produce

Conversation

@YanhuiDua
Copy link
Copy Markdown
Collaborator

@YanhuiDua YanhuiDua commented Feb 26, 2026

  • Sampler 重命名:

Sampler -> _DatasetSampler, SamplerWithReplayBuffer -> Sampler,并且ProduceStrategy默认使用Sampler

  • AsyncProduceStrategy 重命名:AsyncProduceStrategy -> OverProduceStrategy

  • ProduceStrategy 增加两个入参:is_valid_sample_fn, should_continue_fn

  1. is_valid_sample_fn 用于定义判断样本是否有效,过滤的规则写在这里进行传入,默认为 all(sample.status == Status.COMPLETED for sample in samples)
  2. should_continue_fn 传入是否继续生产的条件,默认为completed_samples_count < batch_size

通过增加这两个入参,并且去掉了tqdm进度条,去除ProduceStrategy中hardcode的部分,仅作为一个轻量的控制流,仅通过调用算法/业务相关函数/Module实现对应功能,进度条显示应该有更加优雅的实现方式

…er to Sampler and use Sampler defaultly; 2. add is_valid_sample_fn and should_continue_fn arguments to produce_batch
@YanhuiDua YanhuiDua changed the title Optimizer Sampler && ProduceStrategy Optimizer Sampler and ProduceStrategy Feb 26, 2026
@YanhuiDua YanhuiDua changed the title Optimizer Sampler and ProduceStrategy Optimize the implementation of Sampler and ProduceStrategy Feb 26, 2026
Comment thread xtuner/v1/rl/base/producer.py Outdated
model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
dataloader_cfg: DataloaderConfig
tokenizer: Union[str, PreTrainedTokenizer, PreTrainedTokenizerFast]
prompt_repeat_k: int = 1
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

应该遵循 SamplerCfg.build() 得到对应 Sampler的惯用模式。即DatasetSamplerCfg.build得到DatasetSampler, SamplerCfg.build得到Sampler。
其中静态参数在SamplerCfg中指定,动态参数(比如外部生成的tokenizer, mesh)作为 build()的入参。

Comment thread xtuner/v1/rl/base/agent_loop.py Outdated

class AgentLoopConfig(BaseModel):
model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)
type: Type[AgentLoop] = SingleTurnAgentLoop
Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@jayhenry @hhaAndroid 这里帮忙看下,传AgentLoop的原因是:想要支持用户定义新的AgentLoop后,仅需要修改Config就可以切换AgentLoop。但是没想到更好的传参数方式,ProduceStrategyConfig也是类似的问题

Comment thread xtuner/v1/rl/base/producer.py Outdated
class ProducerConfig(BaseModel):
model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
# TODO: 要考虑怎么方便用户使用自定义的 ProduceStrategy, 并且怎么使用用户需要的参数
type: Type[ProduceStrategy] | Literal["sync", "over"] = "sync"
Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ProducerConfig中也有类似于AgentLoopConfig类似的问题 @jayhenry

Comment thread xtuner/v1/rl/base/producer.py Outdated
# TODO: 要考虑怎么方便用户使用自定义的 ProduceStrategy, 并且怎么使用用户需要的参数
type: Type[ProduceStrategy] | Literal["sync", "over"] = "sync"
staleness_threshold: float = 0.0 # only used for over_produce strategy
is_valid_sample_fn: Optional[Callable[[list[RolloutState]], bool]] = None
Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这两个function放在config里也是为了方便用户仅需修改config就可以完成自定义的功能的组合

@YanhuiDua YanhuiDua merged commit 576cc4a into InternLM:rl_design Feb 27, 2026
0 of 3 checks passed
@YanhuiDua YanhuiDua deleted the optim_produce branch March 3, 2026 11:35
YanhuiDua added a commit that referenced this pull request Apr 27, 2026
* Optimizer Sampler && ProduceStrategy: 1. rename SamplerWithReplayBuffer to Sampler and use Sampler defaultly; 2. add is_valid_sample_fn and should_continue_fn arguments to produce_batch

* Add config for agentloop, agentloopmanager, produceStrategy

* fix ProduceStrategyConfig and AgentLoopConfig
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants