bigscience-workshop / bigscience-workshop/Megatron-DeepSpeed
Calling IndexedDatasetBuilder directly with a best_fit datatype fails
- Dominant language
- Python
- Stars
- 1.4k
- Forks
- 226
- PR merge metrics
- No merged PRs in 30d
Description
I don't know whether this is intended to work or not, but I found the following program:
```
from megatron.data.indexed_dataset import IndexedDatasetBuilder, best_fitting_dtype
best_dtype = best_fitting_dtype(10_000)
IndexedDatasetBuilder("testfile", dtype=best_dtype)
```
leads to an error like:
```
File "/path/to/Megatron-DeepSpeed.git/megatron/data/indexed_dataset.py", line 284, in __init__
self.element_size = self.element_sizes[self.dtype]
KeyError:
```
This shows up because ``best_fitting_dtype`` will return ``numpy.uint16`` for small vocabs:
https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/c680954e0b13232abd5b72711f2032bea1ad65c9/megatron/data/indexed_dataset.py#L25-L27
but that particular type is missing from the ``element_sizes`` table.
https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/c680954e0b13232abd5b72711f2032bea1ad65c9/megatron/data/indexed_dataset.py#L268-L276
which is used in the ``IndexedDatasetBuilder`` constructor here:
https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/c680954e0b13232abd5b72711f2032bea1ad65c9/megatron/data/indexed_dataset.py#L284
Should something like this work?
If so, it seems like either the ``best_fitting_dtype`` function should return a different type or ``uint16`` should be added to the ``element_sizes`` table like https://github.com/bigscience-workshop/Megatron-DeepSpeed/pull/55/commits/f706108acadb2d97dc77f066930accfbdde4942f.
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.