Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

support branch parallel for evoformer #73

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

support branch parallel for evoformer #73

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
62e8185
support branch parallel for evoformer
GuoxiaWang Nov 10, 2022
54a1387
fix conflicts
GuoxiaWang Nov 10, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 26 additions & 0 deletions train_monomer_demo_bp.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
ps -ef | grep "torch" | grep -v grep | awk '{print $2}' | xargs kill -9
ps -ef | grep "unicore-train" | grep -v grep | awk '{print $2}' | xargs kill -9
export MASTER_IP=10.67.228.15
[ -z "${MASTER_IP}" ] && MASTER_IP=127.0.0.1
[ -z "${MASTER_PORT}" ] && MASTER_PORT=12345
[ -z "${n_gpu}" ] && n_gpu=$(nvidia-smi -L | wc -l)
[ -z "${PADDLE_TRAINERS_NUM}" ] && PADDLE_TRAINERS_NUM=1
[ -z "${PADDLE_TRAINER_ID}" ] && PADDLE_TRAINER_ID=0
export NCCL_ASYNC_ERROR_HANDLING=1
export OMP_NUM_THREADS=1
mkdir -p $1
#n_gpu=4
tmp_dir=`mktemp -d`
#model_name=model_init_af2
model_name=model_1_af2
python -m torch.distributed.launch --nproc_per_node=$n_gpu --master_port=$MASTER_PORT --nnodes=$PADDLE_TRAINERS_NUM --node_rank=$PADDLE_TRAINER_ID --master_addr=$MASTER_IP $(which unicore-train) ./example_data/ --user-dir unifold \
--num-workers 8 --ddp-backend=no_c10d \
--task af2 --loss af2 --arch af2 --model-name $model_name \
--optimizer adam --adam-betas '(0.9, 0.999)' --adam-eps 1e-6 --clip-norm 0.0 --per-sample-clip-norm 0.1 --allreduce-fp32-grad \
--lr-scheduler exponential_decay --lr 1e-3 --warmup-updates 1000 --decay-ratio 0.95 --decay-steps 50000 --batch-size 1 \
--update-freq 1 --seed 42 --tensorboard-logdir $1/tsb/ \
--max-update 1000 --max-epoch 1 --log-interval 10 --log-format simple \
--save-interval-updates 500 --validate-interval-updates 500 --keep-interval-updates 40 --no-epoch-checkpoints \
--save-dir $1 --tmp-save-dir $tmp_dir --required-batch-size-multiple 1 --bf16 --ema-decay 0.999 --data-buffer-size 32 --bf16-sr --bp-degree 2
rm -rf $tmp_dir
# --save-dir $1 --tmp-save-dir $tmp_dir --required-batch-size-multiple 1 --ema-decay 0.999 --data-buffer-size 32 --bp-degree 1
4 changes: 2 additions & 2 deletions unifold/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -317,7 +317,7 @@ def base_config():
"pair_dropout": 0.25,
"inf": 1e9,
"eps": 1e-10,
"outer_product_mean_first": False,
"outer_product_mean_pos": 'end',
},
"enabled": True,
},
Expand All @@ -337,7 +337,7 @@ def base_config():
"pair_dropout": 0.25,
"inf": 1e9,
"eps": 1e-10,
"outer_product_mean_first": False,
"outer_product_mean_pos": 'end',
},
"structure_module": {
"d_single": d_single,
Expand Down
24 changes: 12 additions & 12 deletions unifold/data/data_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -619,18 +619,18 @@ def get_pad_size(cur_size, multiplier=4):
return max(multiplier,
((cur_size + multiplier - 1) // multiplier) * multiplier
)
if num_res is not None:
input_num_res = (
protein["aatype"].shape[0]
if "aatype" in protein
else protein["msa_mask"].shape[1]
)
if input_num_res != num_res:
num_res = get_pad_size(input_num_res, 4)
if "extra_msa_mask" in protein:
input_extra_msa_size = protein["extra_msa_mask"].shape[0]
if input_extra_msa_size != extra_msa_size:
extra_msa_size = get_pad_size(input_extra_msa_size, 8)
# if num_res is not None:
# input_num_res = (
# protein["aatype"].shape[0]
# if "aatype" in protein
# else protein["msa_mask"].shape[1]
# )
# if input_num_res != num_res:
# num_res = get_pad_size(input_num_res, 4)
# if "extra_msa_mask" in protein:
# input_extra_msa_size = protein["extra_msa_mask"].shape[0]
# if input_extra_msa_size != extra_msa_size:
# extra_msa_size = get_pad_size(input_extra_msa_size, 8)
pad_size_map = {
N_RES: num_res,
N_MSA: msa_cluster_size,
Expand Down
Loading