zyf981437225-arch/maize-dna-mamba

Maize genomic DNA masked-language model with bidirectional Mamba and cross-layer memory. # Maize DNA-Mamba / OneMaize 本仓库提供玉米基因组 DNA-Mamba 的两阶段预训练流程。当前版本将数据构建、完整性校验、性能基准、分布式 smoke test、正式训练、断点续训和评估拆成可独立执行的步骤，适用于单节点 8×H200，也可先在单卡 A100/H200 上完成小规模验证。 ## 1. 当前训练方案 | 阶段 | 数据 | 序列长度 | 采样方式 | 主要目标 | | --- | --- | ---: | --- | --- | | Phase I | B73，染色体 1–10 | 8,192 bp | 从坐标 0 开始连续、非重叠切片；末尾不足 8,192 bp 的片段保留并 PAD | 建立全基因组基础表示 | | Phase II | B73 + 25 个 NAM 材料 | 16,384 bp | gene / TE / intergenic 按 50% / 30% / 20% 采样，并在候选区域内动态裁剪 | 扩展跨材料和功能区域表示 | 两阶段的定义不可混用：Phase I 不使用 region-aware sampler；Phase II 不使用固定的全基因组非重叠窗口。 当前实现状态： - Phase I 的 B73 manifest 构建、严格校验、dataset、训练配置、benchmark、smoke test 和审计脚本已经完成。 - Phase II 的 16K region-aware sampler 保持原设计；B73 候选区域已经通过长度审计。 - 26 个材料的正式 Phase II metadata 仍需在全部 FASTA/GFF3 到位后构建并验收。 - 8×H200 的最终 batch size、吞吐率和训练预算应以目标计算节点上的 benchmark 结果为准。 ## 2. 模型与训练定义 - backbone：双向 Mamba DNA language model - 层数：24 - hidden size：864 - 参数量：约 121,191,553 - 词表：`A/C/G/T/N`，`PAD=4` - 训练目标：masked language modeling - mask 比例：15% - 被选中位置：80% 替换为 MASK/N，10% 替换为随机碱基，10% 保留原碱基 - 训练集 reverse-complement 概率：0.5 - 验证集和测试集：确定性，不做随机 reverse-complement - 数值精度：BF16 - Phase II 正式材料划分：23 train / 1 validation / 2 test，且 B73 必须在 train split ## 3. 关键入口 | 用途 | 文件 | | --- | --- | | Phase I manifest 构建 | `scripts/build_b73_phase1_8k_manifest.py` | | Phase I manifest 校验 | `scripts/validate_b73_phase1_8k_manifest.py` | | Phase I dataset | `src/dataloaders/datasets/b73_full_genome.py` | | Phase I 训练配置 | `configs/experiment/onemaize_b73_phase1_8k_full_genome.yaml` | | Phase I benchmark | `scripts/benchmark_b73_phase1_8k.py` | | Phase I 启动器 | `scripts/run_onemaize_phase1_h200.sh` | | Phase I Slurm 模板 | `slurm_scripts/run_onemaize_phase1_h200.slurm` | | Phase I 切片统计 | `scripts/stat_b73_8192_full_genome_slicing.py` | | Phase II region metadata 构建 | `scripts/build_onemaize_regions.py` | | Phase II 数据校验 | `scripts/validate_onemaize_data.py` | | Phase II 候选长度审计 | `scripts/audit_onemaize_phase2_candidate_lengths.py` | | Phase II 训练配置 | `configs/experiment/onemaize_b73_phase …