Cifar dataloader does not work
还没有人认领这个 Issue。
评估
- 难度
- 4/5
- 预计耗时
- 3-5 天
- 新手友好度
- 35/100
- Issue 类型
- 缺陷
- 描述清晰度
- 基本清楚
- 活跃度
- 停滞
- 技术栈
- python
- 领域
- data, machine-learning
调研方向
从 algorithms/baselines/self_tuning/jax_nadamw_full_budget.py 开始，重点检查 get_batch_size 和 data_selection，然后使用 submission_runner.py 命令以及上文描述的 CIFAR 工作负载重现该故障。跟踪 batch 如何到达 Flax prefetch_to_device，并验证 CIFAR JAX 运行能够完成，且 shard/device 长度不匹配问题不会出现。
由索引模型根据 Issue 内容生成。
描述
The cifar dataloader no longer works properly with jax algorithms using jax.jit. I did not test to see if pytorch algorithms still work with cifar.
Description
When running jax_nadamw_full_budget.py optimizer with the cifar workload, an error is thrown which says
len(shards) = 128 but len(devices) = 8
Here is the relevant log:
I0918 18:15:28.313870 140360727609984 submission_runner.py:359] Starting training loop.
Traceback (most recent call last):
File "/algorithmic-efficiency/submission_runner.py", line 869, in
app.run(main)
File "/usr/local/lib/python3.11/site-packages/absl/app.py", line 308, in run
_run_main(main, args)
File "/usr/local/lib/python3.11/site-packages/absl/app.py", line 254, in _run_main
sys.exit(main(argv))
^^^^^^^^^^
File "/algorithmic-efficiency/submission_runner.py", line 834, in main
score = score_submission_on_workload(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/algorithmic-efficiency/submission_runner.py", line 747, in score_submission_on_workload
score, _ = train_once(
^^^^^^^^^^^
File "/algorithmic-efficiency/submission_runner.py", line 375, in train_once
batch = data_selection(
^^^^^^^^^^^^^^^
File "/algorithmic-efficiency/algorithms/baselines/self_tuning/jax_nadamw_full_budget.py", line 446, in data_selection
batch = next(input_queue)
^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/site-packages/flax/jax_utils.py", line 147, in prefetch_to_device
enqueue(size) # Fill up the buffer.
^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/site-packages/flax/jax_utils.py", line 145, in enqueue
queue.append(jax.tree_util.tree_map(_prefetch, data))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/site-packages/jax/_src/tree_util.py", line 361, in tree_map
return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/site-packages/jax/_src/tree_util.py", line 361, in
return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
^^^^^^
File "/usr/local/lib/python3.11/site-packages/flax/jax_utils.py", line 141, in _prefetch
return jax.device_put_sharded(list(xs), devices)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/site-packages/jax/_src/api.py", line 2636, in device_put_sharded
raise ValueError(f"len(shards) = {len(shards)} must equal "
ValueError: len(shards) = 128 must equal len(devices) = 8.
2025-09-18 18:15:29.477815: W tensorflow/core/kernels/data/cache_dataset_ops.cc:916] The calling iterator did not fully read the dataset being cached. In order to avoid unexpected truncation of the dataset, the partially cached contents of the dataset will be discarded. This can happen if you have an input pipeline similar todataset.cache().take(k).repeat(). You should usedataset.take(k).cache().repeat()instead.
Steps to Reproduce
- In algorithms/baselines/self_tuning/jax_nadamw_full_budget.py add the following two lines in get_batch_size function:
elif workload_name == 'cifar':
return 128
- Then run the cifar workload in docker:
python submission_runner.py
--framework=jax
--workload=cifar
--experiment_dir=/experiment_runs
--experiment_name=jax_debug_cifar
--data_dir=/data
--tuning_ruleset=self
--submission_path=algorithms/baselines/self_tuning/jax_nadamw_full_budget.py
Source or Possible Fix
I think the cifar is not an officially supported workload, but it can be useful for debugging. So once it is not too much trouble we should fix this.
- 主要语言
- Python
- 星标
- 425
- 派生
- 79
- PR 合并指标
- 30 天内没有已合并 PR
环境准备
- 没有 Dockerfile 或 Docker Compose 文件
- 没有 Pull Request 模板
- 阅读贡献指南
从这里开始
- 先读完整个 Issue，再读项目的贡献指南。
- 在 Issue 下留言说明你要接手 —— 这能避免两个人做同样的事。
- Fork 仓库，在一个分支上完成修改。
- 提交 Pull Request，并在描述里引用这个 Issue 编号。
mlcommons/algorithmic-efficiency 的其他 Issue
-
难度 5/5 一周以上 新手友好度 30/100
-
难度 4/5 3-5 天 新手友好度 38/100
mlcommons/algorithmic-efficiency#924 · 5 条评论 ·
-
难度 5/5 一周以上 新手友好度 25/100
-
难度 5/5 一周以上 新手友好度 20/100
-
Topological overfitting detection: H0 gap catches overfitting before accuracy diverges (r=0.998)未关闭
难度 5/5 一周以上 新手友好度 15/100
查看 mlcommons/algorithmic-efficiency 的全部 Issue
相似的 Issue
-
bug server
难度 2/5 1-3 小时 新手友好度 78/100
维护者通常 1 天内回复
-
难度 2/5 1-3 小时 新手友好度 78/100
sportsdataverse/sportsdataverse-py#641 ·
维护者通常 1 天内回复
-
难度 2/5 1-3 小时 新手友好度 82/100
googleapis/google-cloud-python#18532 ·
维护者通常 1 天内回复
-
难度 2/5 1-3 小时 新手友好度 72/100
维护者通常 1 天内回复
-
难度 2/5 1-3 小时 新手友好度 72/100
维护者通常 1 天内回复