data_param:
    pmd:
        dataset_name: PMD
        add_begin_of_doc_token: True
        add_end_of_doc_token: True
        map_batch_size: 512
        max_num_images: 15
        max_seq_len: 128
        pad_dataset: True
        prob_image_at_end: 0.1
        max_num_samples_per_document: 1
        shuffle_initial_urls_list: True
        shuffle_before_split_by_node_buffer_size: 10000
        shuffle_before_split_by_worker_buffer_size: 10000
        shuffle_after_tarfile_to_samples_buffer_size: 60000
    # laion:
    #     dataset_name: LAION
    #     add_begin_of_doc_token: True
    #     add_end_of_doc_token: True
    #     map_batch_size: 512
    #     max_num_images: 35
    #     max_seq_len: 512
    #     pad_dataset: True
    #     prob_image_at_end: 0.1
    #     max_num_samples_per_document: 1
    #     shuffle_before_split_by_node_buffer_size: 10000
    #     shuffle_before_split_by_worker_buffer_size: 10000
    #     shuffle_after_tarfile_to_samples_buffer_size: 60000
    cm4:
        dataset_name: CM4
        add_begin_of_doc_token: True
        add_end_of_doc_token: True
        map_batch_size: 512
        max_num_images: 15
        max_seq_len: 128
        pad_dataset: True
        p_next: 0.5
        max_num_samples_per_document: 1
        shuffle_initial_urls_list: True
        shuffle_before_split_by_node_buffer_size: 10000
        shuffle_before_split_by_worker_buffer_size: 10000
        shuffle_after_tarfile_to_samples_buffer_size: 12000
    # wiki:
    #     dataset_name: WIKI
    #     add_begin_of_doc_token: True
    #     add_end_of_doc_token: True
    #     map_batch_size: 512
    #     max_num_images: 35
    #     max_seq_len: 512
    #     pad_dataset: True
    #     p_next: 0.5
    #     max_num_samples_per_document: 1
    #     shuffle_before_split_by_node_buffer_size: 10000
    #     shuffle_before_split_by_worker_buffer_size: 10000
    #     shuffle_after_tarfile_to_samples_buffer_size: 12000
    num_workers: 2
    realtime_processing: True
    persistent_workers: False
    proba_interleaving_dataset: [0.15, 0.85]
    use_webdataset: True
hparams:
    tokenizer_name: HuggingFaceM4/huggy-llama-tokenizer-7b
    tokenizer_params: '{"use_fast": True}'
    tokenizer_add_tokens: '[AddedToken("<fake_token_around_image>", rstrip=False, lstrip=False), AddedToken("<image>", rstrip=False, lstrip=False)]'
    tokenizer_add_special_tokens: '{}'
    model_name: huggingface/llama-7b
    model_params:
        vision_embed_dim: 1280
        vision_image_size: 224
        vision_model_name: laion/CLIP-ViT-H-14-laion2B-s32B-b79K
        vision_model_params: '{"id2label":{}, "label2id":{}}'
        tie_word_embeddings: False
        freeze_lm_head: True
        freeze_text_layers: True
        freeze_vision_layers: True
        use_resampler: True
        qk_layer_norms_perceiver: True
        resampler_n_latents: 64
        resampler_depth: 6
        resampler_n_heads: 16
        resampler_head_dim: 96
        alpha_initializer: zeros
        alpha_type: float
        cross_layer_interval: 4
        qk_layer_norms: True
        use_cache: True
    global_batch_size: 8
    batch_size_per_gpu: 4
    gradient_checkpointing: True
    grad_clip: 1.0
    max_num_opt_steps: 2
    seed: 13
    train_logging_activations:
    - jsonl
    train_logging_activations_opt_steps: 250
    train_logging_grad_param_deepspeed:
    - jsonl
    train_logging_grad_param_deepspeed_opt_steps: 250
    train_logging_opt_steps: 1
    train_saving_opt_steps: 1
    do_validation: False
    # kill_switch_path: /fsx/m4/experiments/kill-switch-tr_188_hfc_llama_65_1024_perceiver_new_laion_and_cm4.txt
    wandb_enable: true
    wandb_entity: huggingfacem4
    wandb_log_freq: 100
    wandb_project: VLOOM
    upload_to_s3: False
    timing_break_down: False
optim_param:
    vl_optim: AdamW
    vl_optim_params:
        betas: [0.9, 0.999]
        lr: 0.0001
        weight_decay: 0.1
        no_decay: ["bias", "alpha", "layernorm", "ln", "perceiver_resampler", "layer_norm"]
    vl_lr_scheduler: get_linear_schedule_with_warmup
    vl_lr_scheduler_params:
        last_epoch: -1
        num_warmup_steps: 2_000
        num_training_steps: 500_000
    z_loss: 1e-3