ViperEkura
/
AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix : SFT pipeline position_ids default & doc boundary preservation 

- 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
This commit is contained in:
ViperEkura 2026-07-04 15:59:11 +08:00
parent 8999ca89b8
commit 4e508afa2d
3 changed files with 13 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
assets/docs/preprocessing.md
View File

@ -268,7 +268,7 @@ When `sources` is set, `sections` is ignored.
| `storage_format` | str | `"bin"` | `"bin"` (mmap) or `"h5"` | | `storage_format` | str | `"bin"` | `"bin"` (mmap) or `"h5"` |
| `max_tokens_per_shard` | int | `100000000` | Flush threshold in cumulative tokens | | `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"}`) | | `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"` |
--- ---

2
astrai/config/preprocess_config.py
View File

@ -96,7 +96,7 @@ class OutputConfig(BaseConfig):
storage_format: str = "bin" storage_format: str = "bin"
max_tokens_per_shard: int = 100_000_000 max_tokens_per_shard: int = 100_000_000
dtype: Dict[str, str] = field(default_factory=dict) dtype: Dict[str, str] = field(default_factory=dict)
position_ids_mode: str = "none" position_ids_mode: str = "doc_reset"
@dataclass @dataclass

15
astrai/preprocessing/pipeline.py
View File

@ -144,7 +144,7 @@ class Pipeline:
for key in list(bucket.keys()): for key in list(bucket.keys()):
if key in result: if key in result:
continue continue
bucket[key].append([1] * len(ids)) bucket[key].append([0] * len(ids))
def _iter_items(self): def _iter_items(self):
for path in self.paths: for path in self.paths:
@ -160,6 +160,12 @@ class Pipeline:
idx = shard_idx[domain] idx = shard_idx[domain]
pp = self.config.preprocessing 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) keys = self._packer.apply(dict(keys), pp.max_packed_len, pp.truncation_mode)
tensors: Dict[str, List[torch.Tensor]] = {} tensors: Dict[str, List[torch.Tensor]] = {}
@ -171,9 +177,10 @@ class Pipeline:
torch.tensor(list(chain.from_iterable(ids_list)), dtype=dt) torch.tensor(list(chain.from_iterable(ids_list)), dtype=dt)
] ]
pos_ids = self._position_id.generate(keys.get("sequence", [])) if mode == "continuous" and original_sequences:
if pos_ids: pos_ids = self._position_id.generate(keys.get("sequence", []))
tensors["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)] if pos_ids:
tensors["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)]
self._writer.save(self.output_dir, domain, idx, tensors) self._writer.save(self.output_dir, domain, idx, tensors)
shard_idx[domain] = idx + 1 shard_idx[domain] = idx + 1