From 4e508afa2dba675884d43b3615aad8ee870aa21e Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sat, 4 Jul 2026 15:59:11 +0800 Subject: [PATCH] =?UTF-8?q?=EF=BB=BFfix=20:=20SFT=20pipeline=20position=5F?= =?UTF-8?q?ids=20default=20&=20doc=20boundary=20preservation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - change position_ids_mode default from "none" to "doc_reset" so SFT preprocessing always generates position_ids (was causing dataset load KeyError) - generate per-doc position_ids before packing (doc_reset mode), preserving document boundaries for BFD packing (cross-doc attention leak fix) - change _align_bucket padding from [1] to [0] to avoid accidentally training on loss_mask padding --- assets/docs/preprocessing.md | 2 +- astrai/config/preprocess_config.py | 2 +- astrai/preprocessing/pipeline.py | 15 +++++++++++---- 3 files changed, 13 insertions(+), 6 deletions(-) diff --git a/assets/docs/preprocessing.md b/assets/docs/preprocessing.md index 6246ea4..15afd5c 100644 --- a/assets/docs/preprocessing.md +++ b/assets/docs/preprocessing.md @@ -268,7 +268,7 @@ When `sources` is set, `sections` is ignored. | `storage_format` | str | `"bin"` | `"bin"` (mmap) or `"h5"` | | `max_tokens_per_shard` | int | `100000000` | Flush threshold in cumulative tokens | | `dtype` | dict[str, str] | `{}` | Per-key tensor dtype override (e.g. `{"loss_mask": "bool"}`) | -| `position_ids_mode` | str | `"none"` | How to compute position_ids: `"none"`, `"doc_reset"`, `"continuous"` | +| `position_ids_mode` | str | `"doc_reset"` | How to compute position_ids: `"none"`, `"doc_reset"`, `"continuous"` | --- diff --git a/astrai/config/preprocess_config.py b/astrai/config/preprocess_config.py index 608747f..f10070c 100644 --- a/astrai/config/preprocess_config.py +++ b/astrai/config/preprocess_config.py @@ -96,7 +96,7 @@ class OutputConfig(BaseConfig): storage_format: str = "bin" max_tokens_per_shard: int = 100_000_000 dtype: Dict[str, str] = field(default_factory=dict) - position_ids_mode: str = "none" + position_ids_mode: str = "doc_reset" @dataclass diff --git a/astrai/preprocessing/pipeline.py b/astrai/preprocessing/pipeline.py index e970cb9..5168802 100644 --- a/astrai/preprocessing/pipeline.py +++ b/astrai/preprocessing/pipeline.py @@ -144,7 +144,7 @@ class Pipeline: for key in list(bucket.keys()): if key in result: continue - bucket[key].append([1] * len(ids)) + bucket[key].append([0] * len(ids)) def _iter_items(self): for path in self.paths: @@ -160,6 +160,12 @@ class Pipeline: idx = shard_idx[domain] pp = self.config.preprocessing + original_sequences = keys.get("sequence", []) + mode = self.config.output.position_ids_mode + + if mode == "doc_reset" and original_sequences: + keys["position_ids"] = [list(range(len(s))) for s in original_sequences] + keys = self._packer.apply(dict(keys), pp.max_packed_len, pp.truncation_mode) tensors: Dict[str, List[torch.Tensor]] = {} @@ -171,9 +177,10 @@ class Pipeline: torch.tensor(list(chain.from_iterable(ids_list)), dtype=dt) ] - pos_ids = self._position_id.generate(keys.get("sequence", [])) - if pos_ids: - tensors["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)] + if mode == "continuous" and original_sequences: + pos_ids = self._position_id.generate(keys.get("sequence", [])) + if pos_ids: + tensors["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)] self._writer.save(self.output_dir, domain, idx, tensors) shard_idx[domain] = idx + 1