Skip to content

[bug] Bug with maksed distillation with loss spans #417

@oleksost

Description

@oleksost

🐞 Describe the Bug

When running distillation with loss spans getting this error:


Sample 2 from dataset blended) took 2,825.23 ms to load
2025-12-11 21:57:04,907 [Rank 0] Traceback (most recent call last):
  File "/home/toolkit/dev/Fast-LLM/fast_llm/cli.py", line 27, in fast_llm_main_wrapper
    yield
  File "/home/toolkit/dev/Fast-LLM/fast_llm/cli.py", line 41, in fast_llm_main
    RunnableConfig.parse_and_run(args)
  File "/home/toolkit/dev/Fast-LLM/fast_llm/engine/config_utils/runnable.py", line 27, in parse_and_run
    return cls.get_subclass(args[0]).parse_and_run(args[1:])
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/toolkit/dev/Fast-LLM/fast_llm/engine/config_utils/runnable.py", line 27, in parse_and_run
    return cls.get_subclass(args[0]).parse_and_run(args[1:])
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/toolkit/dev/Fast-LLM/fast_llm/engine/config_utils/runnable.py", line 42, in parse_and_run
    runnable()
  File "/home/toolkit/dev/Fast-LLM/fast_llm/engine/training/config.py", line 388, in runnable
    trainer.run()
  File "/home/toolkit/dev/Fast-LLM/fast_llm/engine/training/trainer.py", line 288, in run
    self._run_training()
  File "/home/toolkit/dev/Fast-LLM/fast_llm/engine/training/trainer.py", line 306, in _run_training
    done, metrics = self._train()
                    ^^^^^^^^^^^^^
  File "/home/toolkit/dev/Fast-LLM/fast_llm/engine/training/trainer.py", line 375, in _train
    reduced_losses, update_successful, train_metrics = self._runner.run_step(
                                                       ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/toolkit/dev/Fast-LLM/fast_llm/engine/schedule/runner.py", line 197, in run_step
    self._train_step(context, step)
  File "/home/toolkit/dev/Fast-LLM/fast_llm/engine/schedule/runner.py", line 311, in _train_step
    output = self._forward(context, step)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/toolkit/dev/Fast-LLM/fast_llm/engine/schedule/runner.py", line 400, in _forward
    self._get_forward_input(context, step),
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/toolkit/dev/Fast-LLM/fast_llm/engine/schedule/runner.py", line 423, in _get_forward_input
    next(context.data_iterator)
  File "/home/toolkit/dev/Fast-LLM/fast_llm/engine/schedule/runner.py", line 327, in _preprocess_data
    micro_batch_data = self._multi_stage.base_model.preprocess_batch(
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/toolkit/dev/Fast-LLM/fast_llm/models/multimodal/model.py", line 188, in preprocess_batch
    preprocessed = super().preprocess_batch(
                   ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/toolkit/dev/Fast-LLM/fast_llm/models/gpt/model.py", line 245, in preprocess_batch
    loss_masking_spans = batch.loss_masking_spans.crop(labels_begin, labels_end)
                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/toolkit/dev/Fast-LLM/fast_llm/data/sample/range.py", line 64, in crop
    return self.__class__([crop_ranges(sample_ranges, begin, end) for sample_ranges in self.ranges], end - begin)
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/toolkit/dev/Fast-LLM/fast_llm/data/sample/range.py", line 21, in crop_ranges
    return [(begin_, end_) for begin_, end_ in cropped_ranges if end_ > begin_]
                                               ^^^^^^^^^^^^^^
  File "/home/toolkit/dev/Fast-LLM/fast_llm/data/sample/range.py", line 20, in <genexpr>
    cropped_ranges = ((max(begin_ - begin, 0), min(end_ - begin, end - begin)) for begin_, end_ in ranges)
                                                                                   ^^^^^^^^^^^^
TypeError: cannot unpack non-iterable numpy.int64 object 

🔄 Steps to Reproduce

Steps to reproduce the behavior:

Set use_loss_masking_spans: yes during distillation.

Metadata

Metadata

Assignees

Labels

bugSomething isn't working

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions