AICB Workload generator

Hi,

I'm trying to use the SimAI AICB workload generator for NS3 simulations, and I have realized that for the same configuration, the generated workload is different, depending on the aiob enablement. Taking this configuration (I modified the scripts to allow TP=8):

sh ./scripts/megatron_workload_with_aiob.sh -m 7 --world_size 16 --tensor_model_parallel_size 8 --pipeline_model_parallel 1 --frame Megatron --global_batch 12 --micro_batch 1 --seq_length 4096 --swiglu --use_flash_attn --aiob_enable --comp_filepath workload/aiob_inputs/Example.txt

If I use aiob without a real GPU (using the aiob input file), the generated workload differs from the one created without the aiob option enabled. For instance, looking at the attention and mlp layers and ignoring the computation times:

With aiob, there two consecutive forward TP operations of 32MB:

mlp_layer -1 2478000 ALLREDUCE 33554432 2478000 NONE 0 2478000 NONE 0 100
attention_layer -1 1820000 ALLREDUCE 33554432 1820000 NONE 0 1820000 NONE 0 100

Without aiob, for each layer there one forward and backward operation of the same size:

attention_layer -1 1 ALLREDUCE 33554432 1 ALLREDUCE 33554432 1 NONE 0 100
mlp_layer -1 1 ALLREDUCE 33554432 1 ALLREDUCE 33554432 1 NONE 0 100

In addition to this difference, there are more. Is there any reason for getting different workloads? Maybe the aiob workload without a real GPU is not reliable?

Thanks

Start with scripts/megatron_workload_with_aiob.sh and the workload/aiob_inputs/Example.txt input used in the reproduction. Generate workloads with and without --aiob_enable using the provided configuration, then trace where attention and MLP operations are emitted. Done means documenting whether the difference is expected without a GPU or correcting the generator so equivalent configurations produce consistent workloads.