OOM on train on 48GPU with batch 1 proc 1

[recoilme]

i waiting sana so long for training on potato, but its not working on A40 with 48GPU(

(sana) root@c88159d783a4:/workspace/sana# bash train_scripts/train.sh   configs/sana_config/1024ms/Sana_1600M_img1024.yaml   --data.data_dir="[asset/example_data]"   --data.type=SanaImgDataset   --model.multi_scale=false
2024-11-26 23:33:11 - [Sana] - INFO - Distributed environment: MULTI_GPU  Backend: nccl
Num processes: 1
Process index: 0
Local process index: 0
Device: cuda:0

Mixed precision type: fp16

2024-11-26 23:33:11 - [Sana] - INFO - Config: 
{
    "data": {
        "data_dir": [
            "asset/example_data"
        ],
        "caption_proportion": {
            "prompt": 1
        },
        "external_caption_suffixes": [
            "",
            "_InternVL2-26B",
            "_VILA1-5-13B"
        ],
        "external_clipscore_suffixes": [
            "_InternVL2-26B_clip_score",
            "_VILA1-5-13B_clip_score",
            "_prompt_clip_score"
        ],
        "clip_thr_temperature": 0.1,
        "clip_thr": 25.0,
        "sort_dataset": false,
        "load_text_feat": false,
        "load_vae_feat": false,
        "transform": "default_train",
        "type": "SanaImgDataset",
        "image_size": 1024,
        "hq_only": false,
        "valid_num": 0,
        "data": null,
        "extra": null
    },
    "model": {
        "model": "SanaMS_1600M_P1_D20",
        "image_size": 1024,
        "mixed_precision": "fp16",
        "fp32_attention": true,
        "load_from": null,
        "resume_from": {
            "checkpoint": "latest",
            "load_ema": false,
            "resume_optimizer": true,
            "resume_lr_scheduler": true
        },
        "aspect_ratio_type": "ASPECT_RATIO_1024",
        "multi_scale": false,
        "pe_interpolation": 1.0,
        "micro_condition": false,
        "attn_type": "linear",
        "autocast_linear_attn": false,
        "ffn_type": "glumbconv",
        "mlp_acts": [
            "silu",
            "silu",
            null
        ],
        "mlp_ratio": 2.5,
        "use_pe": false,
        "qk_norm": false,
        "class_dropout_prob": 0.1,
        "linear_head_dim": 32,
        "cross_norm": false,
        "cfg_scale": 4,
        "guidance_type": "classifier-free",
        "pag_applied_layers": [
            8
        ],
        "extra": null
    },
    "vae": {
        "vae_type": "dc-ae",
        "vae_pretrained": "mit-han-lab/dc-ae-f32c32-sana-1.0",
        "scale_factor": 0.41407,
        "vae_latent_dim": 32,
        "vae_downsample_rate": 32,
        "sample_posterior": true,
        "extra": null
    },
    "text_encoder": {
        "text_encoder_name": "gemma-2-2b-it",
        "caption_channels": 2304,
        "y_norm": true,
        "y_norm_scale_factor": 0.01,
        "model_max_length": 300,
        "chi_prompt": [
            "Given a user prompt, generate an \"Enhanced prompt\" that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:",
            "- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.",
            "- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.",
            "Here are examples of how to transform or refine prompts:",
            "- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.",
            "- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.",
            "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:",
            "User Prompt: "
        ],
        "extra": null
    },
    "scheduler": {
        "train_sampling_steps": 1000,
        "predict_v": true,
        "noise_schedule": "linear_flow",
        "pred_sigma": false,
        "learn_sigma": true,
        "vis_sampler": "flow_dpm-solver",
        "flow_shift": 3.0,
        "weighting_scheme": "logit_normal",
        "logit_mean": 0.0,
        "logit_std": 1.0,
        "extra": null
    },
    "train": {
        "num_workers": 1,
        "seed": 1,
        "train_batch_size": 1,
        "num_epochs": 100,
        "gradient_accumulation_steps": 1,
        "grad_checkpointing": true,
        "gradient_clip": 0.1,
        "gc_step": 1,
        "optimizer": {
            "betas": [
                0.9,
                0.999,
                0.9999
            ],
            "eps": [
                1e-30,
                1e-16
            ],
            "lr": 0.0001,
            "type": "CAMEWrapper",
            "weight_decay": 0.0
        },
        "lr_schedule": "constant",
        "lr_schedule_args": {
            "num_warmup_steps": 2000
        },
        "auto_lr": {
            "rule": "sqrt"
        },
        "ema_rate": 0.9999,
        "eval_batch_size": 16,
        "use_fsdp": false,
        "use_flash_attn": false,
        "eval_sampling_steps": 500,
        "lora_rank": 4,
        "log_interval": 1,
        "mask_type": "null",
        "mask_loss_coef": 0.0,
        "load_mask_index": false,
        "snr_loss": false,
        "real_prompt_ratio": 1.0,
        "save_image_epochs": 1,
        "save_model_epochs": 5,
        "save_model_steps": 500,
        "visualize": true,
        "null_embed_root": "output/pretrained_models/",
        "valid_prompt_embed_root": "output/tmp_embed/",
        "validation_prompts": [
            "dog",
            "portrait photo of a girl, photograph, highly detailed face, depth of field",
            "Self-portrait oil painting, a beautiful cyborg with golden hair, 8k",
            "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
            "A photo of beautiful mountain with realistic sunset and blue lake, highly detailed, masterpiece"
        ],
        "local_save_vis": true,
        "deterministic_validation": true,
        "online_metric": false,
        "eval_metric_step": 2000,
        "online_metric_dir": "metric_helper",
        "work_dir": "output/debug",
        "skip_step": 0,
        "loss_type": "huber",
        "huber_c": 0.001,
        "num_ddim_timesteps": 50,
        "w_max": 15.0,
        "w_min": 3.0,
        "ema_decay": 0.95,
        "debug_nan": false,
        "extra": null
    },
    "work_dir": "output/debug",
    "resume_from": "latest",
    "load_from": null,
    "debug": true,
    "caching": false,
    "report_to": "tensorboard",
    "tracker_project_name": "t2i-evit-baseline",
    "name": "tmp",
    "loss_report_name": "loss"
}
2024-11-26 23:33:11 - [Sana] - INFO - World_size: 1, seed: 1
2024-11-26 23:33:11 - [Sana] - INFO - Initializing: DDP for training
[DC-AE] Loading model from mit-han-lab/dc-ae-f32c32-sana-1.0
Loading checkpoint shards: 100%|███████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00,  5.42it/s]
2024-11-26 23:33:16 - [Sana] - INFO - vae type: dc-ae
2024-11-26 23:33:16 - [Sana] - INFO - Complex Human Instruct: Given a user prompt, generate an "Enhanced prompt" that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:
- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.
- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.
Here are examples of how to transform or refine prompts:
- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.
- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.
Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:
User Prompt: 
2024-11-26 23:33:16 - [Sana] - INFO - v-prediction: True, noise schedule: linear_flow, flow shift: 3.0, flow weighting: logit_normal, logit-mean: 0.0, logit-std: 1.0
2024-11-26 23:33:28 - [Sana] - WARNING - use pe: False, position embed interpolation: 1.0, base size: 32
2024-11-26 23:33:28 - [Sana] - WARNING - attention type: linear; ffn type: glumbconv; autocast linear attn: false
2024-11-26 23:33:41 - [Sana] - INFO - SanaMS:SanaMS_1600M_P1_D20, Model Parameters: 1604.46M
2024-11-26 23:33:41 - [Sana] - INFO - Constructing dataset SanaImgDataset...
2024-11-26 23:33:41 - [Sana] - INFO - Dataset is repeat 2000 times for toy dataset
2024-11-26 23:33:41 - [Sana] - INFO - Dataset samples: 4000
2024-11-26 23:33:41 - [Sana] - INFO - Loading external caption json from: original_filename['', '_InternVL2-26B', '_VILA1-5-13B'].json
2024-11-26 23:33:41 - [Sana] - INFO - Loading external clipscore json from: original_filename['_InternVL2-26B_clip_score', '_VILA1-5-13B_clip_score', '_prompt_clip_score'].json
2024-11-26 23:33:41 - [Sana] - INFO - external caption clipscore threshold: 25.0, temperature: 0.1
2024-11-26 23:33:41 - [Sana] - INFO - T5 max token length: 300
2024-11-26 23:33:41 - [Sana] - INFO - Dataset SanaImgDataset constructed: time: 0.00 s, length (use/ori): 4000/4000
2024-11-26 23:33:41 - [Sana] - INFO - Automatically adapt lr to 0.00001 (using sqrt scaling rule).
2024-11-26 23:33:41 - [Sana] - INFO - CAMEWrapper Optimizer: total 316 param groups, 316 are learnable, 0 are fix. Lr group: 316 params with lr 0.00001; Weight decay group: 316 params with weight decay 0.0.
2024-11-26 23:33:41 - [Sana] - INFO - Lr schedule: constant, num_warmup_steps:2000.
2024-11-26 23:33:41 - [Sana] - WARNING - Basic Setting: lr: 0.00001, bs: 1, gc: True, gc_accum_step: 1, qk norm: False, fp32 attn: True, attn type: linear, ffn type: glumbconv, text encoder: gemma-2-2b-it, captions: {'prompt': 1}, precision: fp16
2024-11-26 23:33:58 - [Sana] - INFO - Epoch: 1 | Global Step: 1 | Local Step: 1 // 4000, total_eta: 71 days, 3:34:13, epoch_eta:17:04:17, time: all:15.368, model:14.333, data:0.201, lm:0.304, vae:0.529, lr:3.125e-09, Cap: VILA1-5-13B, s:(32, 32), loss:4.3361, grad_norm:61.1122
2024-11-26 23:33:58 - [Sana] - INFO - Running validation... 
[rank0]: Traceback (most recent call last):
[rank0]:   File "/workspace/sana/diffusion/model/dpm_solver.py", line 441, in model_fn
[rank0]:     noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2)
[rank0]:   File "/workspace/sana/diffusion/model/dpm_solver.py", line 386, in noise_pred_fn
[rank0]:     output = model(x, t_input, cond, **model_kwargs)
[rank0]:   File "/workspace/sana/diffusion/model/nets/sana_multi_scale.py", line 348, in forward_with_dpmsolver
[rank0]:     model_out = self.forward(x, timestep, y, data_info=data_info, **kwargs)
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/accelerate/utils/operations.py", line 823, in forward
[rank0]:     return model_forward(*args, **kwargs)
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/accelerate/utils/operations.py", line 811, in __call__
[rank0]:     return convert_to_fp32(self.model_forward(*args, **kwargs))
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/amp/autocast_mode.py", line 43, in decorate_autocast
[rank0]:     return func(*args, **kwargs)
[rank0]:   File "/workspace/sana/diffusion/model/nets/sana_multi_scale.py", line 327, in forward
[rank0]:     x = auto_grad_checkpoint(
[rank0]:   File "/workspace/sana/diffusion/model/utils.py", line 72, in auto_grad_checkpoint
[rank0]:     return checkpoint(module, *args, **kwargs)
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/_compile.py", line 31, in inner
[rank0]:     return disable_fn(*args, **kwargs)
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 600, in _fn
[rank0]:     return fn(*args, **kwargs)
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/utils/checkpoint.py", line 481, in checkpoint
[rank0]:     return CheckpointFunction.apply(function, preserve, *args)
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/autograd/function.py", line 574, in apply
[rank0]:     return super().apply(*args, **kwargs)  # type: ignore[misc]
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/utils/checkpoint.py", line 255, in forward
[rank0]:     outputs = run_function(*args)
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:   File "/workspace/sana/diffusion/model/nets/sana_multi_scale.py", line 162, in forward
[rank0]:     x = x + self.drop_path(gate_msa * self.attn(t2i_modulate(self.norm1(x), shift_msa, scale_msa), HW=HW))
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:   File "/workspace/sana/diffusion/model/nets/sana_blocks.py", line 160, in forward
[rank0]:     qkv = self.qkv(x).reshape(B, N, 3, C)
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/nn/modules/linear.py", line 117, in forward
[rank0]:     return F.linear(input, self.weight, self.bias)
[rank0]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 28.00 MiB. GPU 0 has a total capacity of 44.34 GiB of which 20.81 MiB is free. Process 265335 has 44.31 GiB memory in use. Of the allocated memory 42.84 GiB is allocated by PyTorch, and 837.16 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

[rank0]: During handling of the above exception, another exception occurred:

[rank0]: Traceback (most recent call last):
[rank0]:   File "/workspace/sana/train_scripts/train.py", line 974, in <module>
[rank0]:     main()
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/pyrallis/argparsing.py", line 158, in wrapper_inner
[rank0]:     response = fn(cfg, *args, **kwargs)
[rank0]:   File "/workspace/sana/train_scripts/train.py", line 959, in main
[rank0]:     train(
[rank0]:   File "/workspace/sana/train_scripts/train.py", line 479, in train
[rank0]:     log_validation(
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:   File "/workspace/sana/train_scripts/train.py", line 154, in log_validation
[rank0]:     image_logs += run_sampling(init_z=None, label_suffix="", vae=vae, sampler=vis_sampler)
[rank0]:   File "/workspace/sana/train_scripts/train.py", line 127, in run_sampling
[rank0]:     denoised = dpm_solver.sample(
[rank0]:   File "/workspace/sana/diffusion/model/dpm_solver.py", line 1529, in sample
[rank0]:     model_prev_list = [self.model_fn(x, t)]
[rank0]:   File "/workspace/sana/diffusion/model/dpm_solver.py", line 689, in model_fn
[rank0]:     return self.data_prediction_fn(x, t)
[rank0]:   File "/workspace/sana/diffusion/model/dpm_solver.py", line 677, in data_prediction_fn
[rank0]:     noise = self.noise_prediction_fn(x, t)
[rank0]:   File "/workspace/sana/diffusion/model/dpm_solver.py", line 671, in noise_prediction_fn
[rank0]:     return self.model(x, t)
[rank0]:   File "/workspace/sana/diffusion/model/dpm_solver.py", line 616, in <lambda>
[rank0]:     self.model = lambda x, t: model_fn(x, t.expand(x.shape[0]))
[rank0]:   File "/workspace/sana/diffusion/model/dpm_solver.py", line 443, in model_fn
[rank0]:     noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in)[0].chunk(2)
[rank0]:   File "/workspace/sana/diffusion/model/dpm_solver.py", line 386, in noise_pred_fn
[rank0]:     output = model(x, t_input, cond, **model_kwargs)
[rank0]:   File "/workspace/sana/diffusion/model/nets/sana_multi_scale.py", line 348, in forward_with_dpmsolver
[rank0]:     model_out = self.forward(x, timestep, y, data_info=data_info, **kwargs)
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/accelerate/utils/operations.py", line 823, in forward
[rank0]:     return model_forward(*args, **kwargs)
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/accelerate/utils/operations.py", line 811, in __call__
[rank0]:     return convert_to_fp32(self.model_forward(*args, **kwargs))
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/amp/autocast_mode.py", line 43, in decorate_autocast
[rank0]:     return func(*args, **kwargs)
[rank0]:   File "/workspace/sana/diffusion/model/nets/sana_multi_scale.py", line 327, in forward
[rank0]:     x = auto_grad_checkpoint(
[rank0]:   File "/workspace/sana/diffusion/model/utils.py", line 72, in auto_grad_checkpoint
[rank0]:     return checkpoint(module, *args, **kwargs)
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/_compile.py", line 31, in inner
[rank0]:     return disable_fn(*args, **kwargs)
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 600, in _fn
[rank0]:     return fn(*args, **kwargs)
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/utils/checkpoint.py", line 481, in checkpoint
[rank0]:     return CheckpointFunction.apply(function, preserve, *args)
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/autograd/function.py", line 574, in apply
[rank0]:     return super().apply(*args, **kwargs)  # type: ignore[misc]
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/utils/checkpoint.py", line 255, in forward
[rank0]:     outputs = run_function(*args)
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:   File "/workspace/sana/diffusion/model/nets/sana_multi_scale.py", line 162, in forward
[rank0]:     x = x + self.drop_path(gate_msa * self.attn(t2i_modulate(self.norm1(x), shift_msa, scale_msa), HW=HW))
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:   File "/workspace/sana/diffusion/model/nets/sana_blocks.py", line 160, in forward
[rank0]:     qkv = self.qkv(x).reshape(B, N, 3, C)
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:   File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/nn/modules/linear.py", line 117, in forward
[rank0]:     return F.linear(input, self.weight, self.bias)
[rank0]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 28.00 MiB. GPU 0 has a total capacity of 44.34 GiB of which 20.81 MiB is free. Process 265335 has 44.31 GiB memory in use. Of the allocated memory 42.87 GiB is allocated by PyTorch, and 800.33 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
E1126 15:34:00.242000 124914645387072 torch/distributed/elastic/multiprocessing/api.py:833] failed (exitcode: 1) local_rank: 0 (pid: 18768) of binary: /root/miniconda3/envs/sana/bin/python
Traceback (most recent call last):
  File "/root/miniconda3/envs/sana/bin/torchrun", line 8, in <module>
    sys.exit(main())
  File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 348, in wrapper
    return f(*args, **kwargs)
  File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/distributed/run.py", line 901, in main
    run(args)
  File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/distributed/run.py", line 892, in run
    elastic_launch(
  File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/distributed/launcher/api.py", line 133, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
  File "/root/miniconda3/envs/sana/lib/python3.10/site-packages/torch/distributed/launcher/api.py", line 264, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
============================================================
train_scripts/train.py FAILED
------------------------------------------------------------
Failures:
  <NO_OTHER_FAILURES>
------------------------------------------------------------
Root Cause (first observed failure):
[0]:
  time      : 2024-11-26_15:34:00
  host      : c88159d783a4
  rank      : 0 (local_rank: 0)
  exitcode  : 1 (pid: 18768)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
============================================================
(sana) root@c88159d783a4:/workspace/sana# 

[lawrence-cj]

Seems my training is under 48GB

bash train_scripts/train.sh configs/sana_config/1024ms/Sana_1600M_img1024.yaml --data.data_dir="[asset/example_data]" --data.type=SanaImgDataset --model.multi_scale=false --data.load_vae_feat=false --train.train_batch_size=1

refer to:

Actually if you switch optimizer type in config file to AdamW, the GPU memory will be less. We will update a newer Came in the future, which will occupy even less than AdamW.

train:
  optimizer:
    lr: 1.0e-4
    type: AdamW
    weight_decay: 0.01
    eps: 1.0e-8
    betas: [0.9, 0.999]

refer to:

[recoilme]

But how you do it? You have 32x VAE vs 8x in SDXL, less model size and need more then 2.5x against SDXL for train in 1024?

[lawrence-cj]

The model you are using is 1.6B and all the VAE and Text Encoder are all extracting feature online.

[lawrence-cj]

Cool. Hh, I can't imagine the speed. lol : )

[recoilme]

i'm very hope what you will add some minimum optimizations in the future..
its dead for full fine tuning for now, A100 is very expensive
Thx for reply and good model!

[lawrence-cj]

What do you mean it's dead?

[recoilme]

dead

it's dead for full fine-tunings from GPU poor guys
We rent GPU for train. It's very expensive. 48GPU+ for train with batch 1 - its stop factor for most of us
We need latent/TE caching and multi aspect ratio
and probably slow optimizer like adafactor for train fine details like eyes with low LR

[Muinez]

Could you please try this:
#50

[lawrence-cj]

@Muinez , nice man. I'll try this out.

[recoilme]

@Muinez thank you! bucketing works nice, code is clean and very simple!

i add checkpoints dir creation around line 500 (it fail on 1st run):

checkpoints_dir = osp.join(config.work_dir, "checkpoints")
os.makedirs(checkpoints_dir, exist_ok=True)

its not work for me on 48Gb/A40 with batch 1 with CAME optimizer (default settings, exclude batch/num proc) - OOM after generating validation images with wandb (probably after decoding - you may see memory spikes at validation)

but its looks like it do something strange and work on Adam https://wandb.ai/recoilme/sana/runs/1?nw=nwuserrecoilme

Looks like i train from clean checkpoint, something wrong?

What do you think about add adafactor? Is it possible? I have totally better results with this optimizer on SDXL. Adam is like tank - good only in pretrain case..

[lawrence-cj]

The loss means you are training from scratch. (no ckpts loaded)